From 4d7029cf9b16b4c9cb7c91bf6254836f89ef0559 Mon Sep 17 00:00:00 2001
From: Michael Bang <mic@vbang.dk>
Date: Sun, 19 Nov 2023 13:43:57 +0100
Subject: [PATCH] datatables cs2: rewrite player equipment tracking

NOTE: this is not 100% accurate yet, but it's a big step in the right
direction.

Thanks to esbengc for describing the required logic!
---
 pkg/demoinfocs/datatables.go | 80 +++++++++++++++++++++++-------------
 1 file changed, 51 insertions(+), 29 deletions(-)

diff --git a/pkg/demoinfocs/datatables.go b/pkg/demoinfocs/datatables.go
index b73a6b30..a1df0858 100644
--- a/pkg/demoinfocs/datatables.go
+++ b/pkg/demoinfocs/datatables.go
@@ -688,44 +688,66 @@ 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
+	playerInventory := make(map[int]*common.Equipment)
+
+	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)]
 
-				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
-				}
+			entityID, wep := getWep(val)
+			wep.Owner = pl
 
-				// Clear previous owner
-				if wep.Owner != nil && wep.Entity != nil {
-					delete(wep.Owner.Inventory, wep.Entity.ID())
-				}
+			entityWasCreated := entityID != constants.EntityHandleIndexMaskSource2
 
-				// 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
+			if uint64(i) < inventorySize {
+				if entityWasCreated {
+					existingWeapon, exists := playerInventory[i]
+					if exists {
+						delete(pl.Inventory, existingWeapon.Entity.ID())
+					}
+					pl.Inventory[int(entityID)] = wep
+					playerInventory[i] = wep
+				} else {
+					delete(pl.Inventory, int(entityID))
 				}
-				delete(pl.Inventory, int(cache[i2]))
 
-				cache[i2] = 0
+				setPlayerInventory()
 			}
 		}
 		property := pawnEntity.Property(playerWeaponPrefixS2 + fmt.Sprintf("%04d", i))