feat(texture): dispose of textures when no longer in use

Author: fallenoak

src/lib/index.ts

@@ -2,7 +2,7 @@ export * from './util.js';
 export * from './controls/OrbitControls.js';
 export * from './AssetManager.js';
 export * from './FormatManager.js';
-export * from './TextureManager.js';
+export * from './texture/TextureManager.js';
 export * from './SceneWorker.js';
 export * from './terrain/TerrainManager.js';
 export * from './terrain/TerrainMesh.js';

src/lib/terrain/TerrainManager.ts

@@ -2,7 +2,7 @@ import * as THREE from 'three';
 import { MapChunk, MapArea } from '@wowserhq/format';
 import { createSplatTexture } from './material.js';
 import { createTerrainIndexBuffer, createTerrainVertexBuffer } from './geometry.js';
-import TextureManager from '../TextureManager.js';
+import TextureManager from '../texture/TextureManager.js';
 import TerrainMaterial from './TerrainMaterial.js';
 import TerrainMesh from './TerrainMesh.js';
 

src/lib/texture/ManagedCompressedTexture.ts

@@ -0,0 +1,49 @@
+import * as THREE from 'three';
+import TextureManager from './TextureManager.js';
+
+class ManagedCompressedTexture extends THREE.CompressedTexture {
+  #manager: TextureManager;
+  #refId: string;
+
+  constructor(
+    manager: TextureManager,
+    refId: string,
+    mipmaps: ImageData[],
+    width: number,
+    height: number,
+    format: THREE.CompressedPixelFormat,
+    type?: THREE.TextureDataType,
+    mapping?: THREE.Mapping,
+    wrapS?: THREE.Wrapping,
+    wrapT?: THREE.Wrapping,
+    magFilter?: THREE.MagnificationTextureFilter,
+    minFilter?: THREE.MinificationTextureFilter,
+    anisotropy?: number,
+    colorSpace?: THREE.ColorSpace,
+  ) {
+    super(
+      mipmaps,
+      width,
+      height,
+      format,
+      type,
+      mapping,
+      wrapS,
+      wrapT,
+      magFilter,
+      minFilter,
+      anisotropy,
+      colorSpace,
+    );
+
+    this.#manager = manager;
+    this.#refId = refId;
+  }
+
+  dispose() {
+    this.#manager.deref(this.#refId);
+  }
+}
+
+export default ManagedCompressedTexture;
+export { ManagedCompressedTexture };

src/lib/texture/ManagedDataTexture.ts

@@ -0,0 +1,49 @@
+import * as THREE from 'three';
+import TextureManager from './TextureManager.js';
+
+class ManagedDataTexture extends THREE.DataTexture {
+  #manager: TextureManager;
+  #refId: string;
+
+  constructor(
+    manager: TextureManager,
+    refId: string,
+    data?: BufferSource | null,
+    width?: number,
+    height?: number,
+    format?: THREE.PixelFormat,
+    type?: THREE.TextureDataType,
+    mapping?: THREE.Mapping,
+    wrapS?: THREE.Wrapping,
+    wrapT?: THREE.Wrapping,
+    magFilter?: THREE.MagnificationTextureFilter,
+    minFilter?: THREE.MinificationTextureFilter,
+    anisotropy?: number,
+    colorSpace?: THREE.ColorSpace,
+  ) {
+    super(
+      data,
+      width,
+      height,
+      format,
+      type,
+      mapping,
+      wrapS,
+      wrapT,
+      magFilter,
+      minFilter,
+      anisotropy,
+      colorSpace,
+    );
+
+    this.#manager = manager;
+    this.#refId = refId;
+  }
+
+  dispose() {
+    this.#manager.deref(this.#refId);
+  }
+}
+
+export default ManagedDataTexture;
+export { ManagedDataTexture };

src/lib/TextureManager.ts renamed to src/lib/texture/TextureManager.ts

@@ -1,7 +1,9 @@
 import * as THREE from 'three';
 import { Blp, BLP_IMAGE_FORMAT } from '@wowserhq/format';
-import FormatManager from './FormatManager.js';
-import { normalizePath } from './util.js';
+import ManagedCompressedTexture from './ManagedCompressedTexture.js';
+import ManagedDataTexture from './ManagedDataTexture.js';
+import FormatManager from '../FormatManager.js';
+import { normalizePath } from '../util.js';
 
 const THREE_TEXTURE_FORMAT: Record<number, THREE.PixelFormat | THREE.CompressedPixelFormat> = {
   [BLP_IMAGE_FORMAT.IMAGE_DXT1]: THREE.RGBA_S3TC_DXT1_Format,
@@ -14,6 +16,7 @@ class TextureManager {
   #formatManager: FormatManager;
   #loaded = new Map<string, THREE.Texture>();
   #loading = new Map<string, Promise<THREE.Texture>>();
+  #refs = new Map<string, number>();
 
   constructor(formatManager: FormatManager) {
     this.#formatManager = formatManager;
@@ -26,26 +29,64 @@ class TextureManager {
     minFilter: THREE.MinificationTextureFilter = THREE.LinearMipmapLinearFilter,
     magFilter: THREE.MagnificationTextureFilter = THREE.LinearFilter,
   ) {
-    const cacheKey = [normalizePath(path), wrapS, wrapT, minFilter, magFilter].join(':');
+    const refId = [normalizePath(path), wrapS, wrapT, minFilter, magFilter].join(':');
+    this.#ref(refId);
 
-    const loaded = this.#loaded.get(cacheKey);
+    const loaded = this.#loaded.get(refId);
     if (loaded) {
       return Promise.resolve(loaded);
     }
 
-    const alreadyLoading = this.#loading.get(cacheKey);
+    const alreadyLoading = this.#loading.get(refId);
     if (alreadyLoading) {
       return alreadyLoading;
     }
 
-    const loading = this.#load(cacheKey, path, wrapS, wrapT, minFilter, magFilter);
-    this.#loading.set(cacheKey, loading);
+    const loading = this.#load(refId, path, wrapS, wrapT, minFilter, magFilter);
+    this.#loading.set(refId, loading);
 
     return loading;
   }
 
+  deref(refId: string) {
+    let refCount = this.#refs.get(refId);
+
+    // Unknown ref
+
+    if (refCount === undefined) {
+      return;
+    }
+
+    // Decrement
+
+    refCount--;
+
+    if (refCount > 0) {
+      this.#refs.set(refId, refCount);
+      return;
+    }
+
+    // Dispose
+
+    const texture = this.#loaded.get(refId);
+    if (texture) {
+      this.#loaded.delete(refId);
+      texture.dispose();
+    }
+
+    this.#refs.delete(refId);
+  }
+
+  #ref(refId: string) {
+    let refCount = this.#refs.get(refId) || 0;
+
+    refCount++;
+
+    this.#refs.set(refId, refCount);
+  }
+
   async #load(
-    cacheKey: string,
+    refId: string,
     path: string,
     wrapS: THREE.Wrapping,
     wrapT: THREE.Wrapping,
@@ -56,7 +97,7 @@ class TextureManager {
     try {
       blp = await this.#formatManager.get(path, Blp);
     } catch (error) {
-      this.#loading.delete(cacheKey);
+      this.#loading.delete(refId);
       throw error;
     }
 
@@ -66,24 +107,28 @@ class TextureManager {
 
     const threeFormat = THREE_TEXTURE_FORMAT[imageFormat];
     if (threeFormat === undefined) {
-      this.#loading.delete(cacheKey);
+      this.#loading.delete(refId);
       throw new Error(`Unsupported texture format: ${imageFormat}`);
     }
 
-    let texture: THREE.CompressedTexture | THREE.DataTexture;
+    let texture: ManagedCompressedTexture | ManagedDataTexture;
     if (
       imageFormat === BLP_IMAGE_FORMAT.IMAGE_DXT1 ||
       imageFormat === BLP_IMAGE_FORMAT.IMAGE_DXT3 ||
       imageFormat === BLP_IMAGE_FORMAT.IMAGE_DXT5
     ) {
-      texture = new THREE.CompressedTexture(
+      texture = new ManagedCompressedTexture(
+        this,
+        refId,
         null,
         blp.width,
         blp.height,
         threeFormat as THREE.CompressedPixelFormat,
       );
     } else if (imageFormat === BLP_IMAGE_FORMAT.IMAGE_ABGR8888) {
-      texture = new THREE.DataTexture(
+      texture = new ManagedDataTexture(
+        this,
+        refId,
         null,
         blp.width,
         blp.height,
@@ -98,18 +143,16 @@ class TextureManager {
     texture.magFilter = magFilter;
     texture.anisotropy = 16;
     texture.name = normalizePath(path).split('/').at(-1);
-    texture.userData.cacheKey = cacheKey;
 
     // All newly loaded textures need to be flagged for upload to the GPU
     texture.needsUpdate = true;
 
-    this.#loaded.set(cacheKey, texture);
-    this.#loading.delete(cacheKey);
+    this.#loaded.set(refId, texture);
+    this.#loading.delete(refId);
 
     return texture;
   }
 }
 
 export default TextureManager;
-
 export { TextureManager };