From a4d7aaccdc8555d08babbd0256baf361296dc10e Mon Sep 17 00:00:00 2001
From: Uladzislau Nikalayevich <thenormalnij@gmail.com>
Date: Tue, 27 Aug 2024 22:20:43 +0300
Subject: [PATCH] Better checks

---
 Client/game_sa/CGameSA.cpp        | 22 +++++++++++++++++-----
 Client/game_sa/CModelInfoSA.cpp   |  2 ++
 Client/game_sa/CPoolSAInterface.h |  2 +-
 Client/game_sa/CTxdPoolSA.cpp     |  9 +++++++++
 Client/game_sa/CTxdPoolSA.h       |  1 +
 Client/sdk/game/CTxdPool.h        |  1 +
 6 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/Client/game_sa/CGameSA.cpp b/Client/game_sa/CGameSA.cpp
index 7222c8db1d..482a283ac0 100644
--- a/Client/game_sa/CGameSA.cpp
+++ b/Client/game_sa/CGameSA.cpp
@@ -1051,17 +1051,29 @@ bool CGameSA::SetBuildingPoolSize(size_t size)
 
 void CGameSA::UnloadUnusedModels()
 {
-    for (std::size_t id = 0; id < GetBaseIDforCOL(); id++)
+    // Unload DFF's
+    // CJ should not be unloaded
+    const std::size_t baseIdForTxd = GetBaseIDforTXD();
+    for (std::size_t id = 1; id < baseIdForTxd; id++)
     {
         CStreamingInfo* streamingInfo = m_pStreaming->GetStreamingInfo(id);
         if (streamingInfo->loadState != eModelLoadState::LOADSTATE_NOT_LOADED && streamingInfo->sizeInBlocks > 0)
         {
-            if (ModelInfo[id].GetRefCount() == 0)
-            {
-                m_pStreaming->RemoveModel(id);
-            }
+            CModelInfoSA& model = ModelInfo[id];
+            if (model.GetRefCount() == 0)
+                model.UnloadUnused();
         };
     }
+    // Unload TXD
+    for (std::size_t id = baseIdForTxd; id < GetBaseIDforCOL(); id++)
+    {
+        CStreamingInfo* streamingInfo = m_pStreaming->GetStreamingInfo(id);
+        std::size_t     refsCount = GetPools()->GetTxdPool().GetRefsCount(id - baseIdForTxd);
+        if (streamingInfo->loadState != eModelLoadState::LOADSTATE_NOT_LOADED && streamingInfo->sizeInBlocks > 0 && refsCount == 0)
+        {
+            GetStreaming()->RemoveModel(id);
+        }
+    }
 }
 
 // Ensure models have the default lod distances
diff --git a/Client/game_sa/CModelInfoSA.cpp b/Client/game_sa/CModelInfoSA.cpp
index cf775a55ee..72b2c92b74 100644
--- a/Client/game_sa/CModelInfoSA.cpp
+++ b/Client/game_sa/CModelInfoSA.cpp
@@ -442,6 +442,8 @@ void CModelInfoSA::Remove()
 
 bool CModelInfoSA::UnloadUnused()
 {
+    m_pInterface = ppModelInfo[m_dwModelID];
+
     if (m_pInterface->usNumberOfRefs == 0 && !m_pCustomClump && !m_pCustomColModel)
     {
         pGame->GetStreaming()->RemoveModel(m_dwModelID);
diff --git a/Client/game_sa/CPoolSAInterface.h b/Client/game_sa/CPoolSAInterface.h
index 9f48bc63d3..f879775b3b 100644
--- a/Client/game_sa/CPoolSAInterface.h
+++ b/Client/game_sa/CPoolSAInterface.h
@@ -122,7 +122,7 @@ class CPoolSAInterface
         return !IsEmpty(index);
     }
 
-    B* GetObject(std::int32_t objectIndex) { return &m_pObjects[objectIndex]; }
+    B* GetObject(std::int32_t objectIndex) const { return &m_pObjects[objectIndex]; }
 
     uint GetObjectIndex(B* pObject) { return ((DWORD)pObject - (DWORD)m_pObjects) / sizeof(B); }
 
diff --git a/Client/game_sa/CTxdPoolSA.cpp b/Client/game_sa/CTxdPoolSA.cpp
index 248f825836..36be04fce5 100644
--- a/Client/game_sa/CTxdPoolSA.cpp
+++ b/Client/game_sa/CTxdPoolSA.cpp
@@ -55,3 +55,12 @@ std::uint16_t CTxdPoolSA::GetFreeTextureDictonarySlot()
 {
     return (*m_ppTxdPoolInterface)->GetFreeSlot();
 }
+
+std::uint16_t CTxdPoolSA::GetRefsCount(std::uint16_t slot) const
+{
+    CTextureDictonarySAInterface* pTxd = (*m_ppTxdPoolInterface)->GetObject(slot);
+    if (!pTxd)
+        return -1;
+
+    return pTxd->usUsagesCount;
+}
diff --git a/Client/game_sa/CTxdPoolSA.h b/Client/game_sa/CTxdPoolSA.h
index c1f0ed8ece..daaf3a1ccb 100644
--- a/Client/game_sa/CTxdPoolSA.h
+++ b/Client/game_sa/CTxdPoolSA.h
@@ -25,6 +25,7 @@ class CTxdPoolSA final : public CTxdPool
     bool          IsFreeTextureDictonarySlot(std::uint32_t uiTxdId);
 
     std::uint16_t GetFreeTextureDictonarySlot();
+    std::uint16_t GetRefsCount(std::uint16_t slot) const;
 
 private:
     CPoolSAInterface<CTextureDictonarySAInterface>** m_ppTxdPoolInterface;
diff --git a/Client/sdk/game/CTxdPool.h b/Client/sdk/game/CTxdPool.h
index 1e4ff3493b..f186d8dff9 100644
--- a/Client/sdk/game/CTxdPool.h
+++ b/Client/sdk/game/CTxdPool.h
@@ -19,4 +19,5 @@ class CTxdPool
     virtual bool          IsFreeTextureDictonarySlot(std::uint32_t uiTxdID) = 0;
 
     virtual std::uint16_t GetFreeTextureDictonarySlot() = 0;
+    virtual std::uint16_t GetRefsCount(std::uint16_t slot) const = 0;
 };