dispose models on unload

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"),
+    ),
+  };
 }

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]);
+}

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,
     ),
   };
 }