From 100c7391fbb81c79d897375bb4d13c4908b2912b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tr=C3=A8sAbhi?= <43380238+tresabhi@users.noreply.github.com>
Date: Tue, 12 May 2026 17:54:50 -0500
Subject: [PATCH] dispose models on unload

---
 packages/website/src/hooks/useArmor.ts   | 34 ++++++------
 packages/website/src/hooks/useDispose.ts | 67 ++++++++++++++++++++++++
 packages/website/src/hooks/useModel.ts   | 12 +++--
 3 files changed, 90 insertions(+), 23 deletions(-)
 create mode 100644 packages/website/src/hooks/useDispose.ts

diff --git a/packages/website/src/hooks/useArmor.ts b/packages/website/src/hooks/useArmor.ts
index e1a5f05a6..d697bd0ba 100644
--- a/packages/website/src/hooks/useArmor.ts
+++ b/packages/website/src/hooks/useArmor.ts
@@ -1,23 +1,19 @@
-import { asset } from '@blitzkit/core';
-import { type ObjectMap, useLoader } from '@react-three/fiber';
-import { type GLTF, GLTFLoader } from 'three-stdlib';
-
-const cache: Record<
-  number,
-  { gltf: GLTF & ObjectMap; hasDynamicArmor: boolean }
-> = {};
+import { asset } from "@blitzkit/core";
+import { useLoader } from "@react-three/fiber";
+import { GLTFLoader } from "three-stdlib";
+import { useDispose } from "./useDispose";
 
 export function useArmor(id: number) {
-  const gltf = useLoader(GLTFLoader, asset(`3d/tanks/armor/${id}.glb`));
+  const path = asset(`3d/tanks/armor/${id}.glb`);
+  const gltf = useLoader(GLTFLoader, path);
 
-  if (!cache[id]) {
-    cache[id] = {
-      gltf,
-      hasDynamicArmor: Object.values(gltf.nodes).some(
-        (node) =>
-          node.name.includes('state_00') || node.name.includes('state_01'),
-      ),
-    };
-  }
-  return cache[id];
+  useDispose(gltf, path);
+
+  return {
+    gltf,
+    hasDynamicArmor: Object.values(gltf.nodes).some(
+      (node) =>
+        node.name.includes("state_00") || node.name.includes("state_01"),
+    ),
+  };
 }
diff --git a/packages/website/src/hooks/useDispose.ts b/packages/website/src/hooks/useDispose.ts
new file mode 100644
index 000000000..dffe2dc81
--- /dev/null
+++ b/packages/website/src/hooks/useDispose.ts
@@ -0,0 +1,67 @@
+import { useLoader } from "@react-three/fiber";
+import { useEffect } from "react";
+import { Cache } from "three";
+import { GLTFLoader } from "three-stdlib";
+
+export function useDispose(gltf: any, path: string) {
+  useEffect(() => {
+    return () => {
+      const disposedGeometries = new Set<string>();
+      const disposedMaterials = new Set<string>();
+      const disposedTextures = new Set<string>();
+
+      gltf.scene.traverse((obj: any) => {
+        if (obj.geometry) {
+          const geometry = obj.geometry;
+
+          if (!disposedGeometries.has(geometry.uuid)) {
+            disposedGeometries.add(geometry.uuid);
+            geometry.dispose();
+          }
+        }
+
+        if (obj.material) {
+          const materials = Array.isArray(obj.material)
+            ? obj.material
+            : [obj.material];
+
+          for (const material of materials) {
+            if (!disposedMaterials.has(material.uuid)) {
+              disposedMaterials.add(material.uuid);
+
+              for (const key in material) {
+                const value = material[key];
+
+                if (
+                  value &&
+                  value.isTexture &&
+                  !disposedTextures.has(value.uuid)
+                ) {
+                  disposedTextures.add(value.uuid);
+                }
+              }
+
+              for (const key in material) {
+                const value = material[key];
+
+                if (value && value.isTexture) {
+                  value.source.data.close();
+                  value.dispose();
+                }
+              }
+
+              material.dispose();
+            }
+          }
+        }
+
+        if (obj.skeleton) {
+          obj.skeleton.dispose?.();
+        }
+      });
+
+      useLoader.clear(GLTFLoader, path);
+      Cache.remove(path);
+    };
+  }, [gltf, path]);
+}
diff --git a/packages/website/src/hooks/useModel.ts b/packages/website/src/hooks/useModel.ts
index 83bfc7cfa..c8d532bd5 100644
--- a/packages/website/src/hooks/useModel.ts
+++ b/packages/website/src/hooks/useModel.ts
@@ -1,10 +1,14 @@
 import { asset } from "@blitzkit/core";
-import { type ObjectMap, useLoader } from "@react-three/fiber";
+import { useLoader } from "@react-three/fiber";
 import { Mesh, MeshStandardMaterial } from "three";
-import { type GLTF, GLTFLoader } from "three-stdlib";
+import { GLTFLoader } from "three-stdlib";
+import { useDispose } from "./useDispose";
 
 export function useModel(id: number) {
-  const gltf = useLoader(GLTFLoader, asset(`3d/tanks/models/${id}.glb`));
+  const path = asset(`3d/tanks/models/${id}.glb`);
+  const gltf = useLoader(GLTFLoader, path);
+
+  useDispose(gltf, path);
 
   return {
     gltf,
@@ -12,7 +16,7 @@ export function useModel(id: number) {
       (node) =>
         node instanceof Mesh &&
         node.material instanceof MeshStandardMaterial &&
-        node.material.roughnessMap !== null
+        node.material.roughnessMap !== null,
     ),
   };
 }