Merge pull request #478 from micvbang/micvbang/simplify-weapon-inventory

datatables cs2: rewrite player equipment tracking

Author: markus-wa

pkg/demoinfocs/datatables.go

@@ -688,46 +688,83 @@ func (p *parser) bindPlayerWeapons(playerEntity st.Entity, pl *common.Player) {
 }
 
 func (p *parser) bindPlayerWeaponsS2(pawnEntity st.Entity, pl *common.Player) {
-	var cache [maxWeapons]uint64
-	for i := range cache {
-		i2 := i // Copy for passing to handler
+	const inventoryCapacity = 64
+
+	var inventorySize uint64 = 64
+
+	type eq struct {
+		*common.Equipment
+		entityID int
+	}
+
+	playerInventory := make(map[int]eq)
+
+	getWep := func(wepSlotPropertyValue st.PropertyValue) (uint64, *common.Equipment) {
+		entityID := wepSlotPropertyValue.S2UInt64() & constants.EntityHandleIndexMaskSource2
+		wep := p.gameState.weapons[int(entityID)]
+
+		if wep == nil {
+			// sometimes a weapon is assigned to a player before the weapon entity is created
+			wep = common.NewEquipment(common.EqUnknown)
+			p.gameState.weapons[int(entityID)] = wep
+		}
+
+		return entityID, wep
+	}
+
+	setPlayerInventory := func() {
+		inventory := make(map[int]*common.Equipment, inventorySize)
+
+		for i := uint64(0); i < inventorySize; i++ {
+			val := pawnEntity.Property(playerWeaponPrefixS2 + fmt.Sprintf("%04d", i)).Value()
+			if val.Any == nil {
+				continue
+			}
+
+			entityID, wep := getWep(val)
+			inventory[int(entityID)] = wep
+		}
+
+		pl.Inventory = inventory
+	}
+
+	pawnEntity.Property("m_pWeaponServices.m_hMyWeapons").OnUpdate(func(pv st.PropertyValue) {
+		inventorySize = pv.S2UInt64()
+		setPlayerInventory()
+	})
+
+	for i := 0; i < inventoryCapacity; i++ {
+		i := i
 		updateWeapon := func(val st.PropertyValue) {
 			if val.Any == nil {
 				return
 			}
-			entityID := val.S2UInt64() & constants.EntityHandleIndexMaskSource2
-			if entityID != constants.EntityHandleIndexMaskSource2 {
-				if cache[i2] != 0 {
-					// Player already has a weapon in this slot.
-					delete(pl.Inventory, int(cache[i2]))
-				}
-				cache[i2] = entityID
 
-				wep := p.gameState.weapons[int(entityID)]
+			entityID, wep := getWep(val)
+			wep.Owner = pl
 
-				if wep == nil {
-					// sometimes a weapon is assigned to a player before the weapon entity is created
-					wep = common.NewEquipment(common.EqUnknown)
-					p.gameState.weapons[int(entityID)] = wep
-				}
+			entityWasCreated := entityID != constants.EntityHandleIndexMaskSource2
 
-				// Clear previous owner
-				if wep.Owner != nil && wep.Entity != nil {
-					delete(wep.Owner.Inventory, wep.Entity.ID())
-				}
+			if uint64(i) < inventorySize {
+				if entityWasCreated {
+					existingWeapon, exists := playerInventory[i]
+					if exists {
+						delete(pl.Inventory, existingWeapon.entityID)
+					}
 
-				// Attribute weapon to player
-				wep.Owner = pl
-				pl.Inventory[int(entityID)] = wep
-			} else {
-				if cache[i2] != 0 && pl.Inventory[int(cache[i2])] != nil {
-					pl.Inventory[int(cache[i2])].Owner = nil
+					pl.Inventory[int(entityID)] = wep
+					playerInventory[i] = eq{
+						Equipment: wep,
+						entityID:  int(entityID),
+					}
+				} else {
+					delete(pl.Inventory, int(entityID))
 				}
-				delete(pl.Inventory, int(cache[i2]))
 
-				cache[i2] = 0
+				setPlayerInventory()
 			}
 		}
+
 		property := pawnEntity.Property(playerWeaponPrefixS2 + fmt.Sprintf("%04d", i))
 		updateWeapon(property.Value())
 		property.OnUpdate(updateWeapon)