Fix memory leak in ReadbackSplatSorter

Author: mrxz

examples/hello-world/index.html

@@ -23,7 +23,7 @@
   </script>
   <script type="module">
     import * as THREE from "three";
-    import { SplatLoader, Splat } from "@sparkjsdev/spark";
+    import { SplatLoader } from "@sparkjsdev/spark";
     import { getAssetFileURL } from "/examples/js/get-asset-url.js";
 
     const scene = new THREE.Scene();

examples/interactivity/index.html

@@ -175,7 +175,6 @@ <h2>Food scans by <a href="https://x.com/tipatat" target="_blank">Tipatat</a></h
           if (food) food.opacity = 1 - easeInOutSine(fadeOutTime / TRANSITION_LENGTH);
         } else {
           // Fade out finished
-          if (food) food.dispose();
           fadeOutTime = null;
           // Fade in next food
           fadeInTime = 0;

src/BatchedSplat.ts

@@ -199,7 +199,7 @@ class BatchedSplatData implements SplatData {
   getInstanceIndexFor(splatIndex: number): number {
     let instanceIndex = 0;
     let instanceEnd = this.sources[instanceIndex].numSplats;
-    while (splatIndex > instanceEnd) {
+    while (splatIndex >= instanceEnd) {
       instanceIndex++;
       instanceEnd += this.sources[instanceIndex].numSplats;
     }

src/Splat.ts

@@ -562,7 +562,9 @@ export class Splat extends THREE.Mesh<SplatGeometry, THREE.ShaderMaterial> {
   }
 
   dispose() {
-    // TODO
+    this.geometry.dispose();
+    this.material.dispose();
+    this.splatData.dispose();
   }
 
   set needsUpdate(value: boolean) {

src/SplatSorter.ts

@@ -70,7 +70,9 @@ export class ReadbackSplatSorter implements SplatSorter {
 
   private capacity = 0;
   private target?: THREE.WebGLArrayRenderTarget;
-  private readbackBufferPool: ReadbackBufferPool;
+  private readonly readbackBufferPool: ReadbackBufferPool;
+
+  private readonly material: THREE.RawShaderMaterial;
 
   constructor(options: ReadbackSorterOptions = {}) {
     this.sortRadial = options.sortRadial ?? false;
@@ -79,6 +81,7 @@ export class ReadbackSplatSorter implements SplatSorter {
     this.sort32 = options.sort32 ?? false;
 
     this.readbackBufferPool = new ReadbackBufferPool();
+    this.material = ReadbackSplatSorter.createMaterial();
   }
 
   async sort(
@@ -92,16 +95,6 @@ export class ReadbackSplatSorter implements SplatSorter {
     const maxSplats = splatData.maxSplats;
     const numSplats = splatData.numSplats;
 
-    const material = this.createMaterial(splatData);
-    material.uniforms.sortRadial.value = this.sortRadial;
-    material.uniforms.sortDepthBias.value = this.depthBias;
-    material.uniforms.sort360.value = this.sort360;
-    material.uniforms.splatModelViewMatrix.value.multiplyMatrices(
-      camera.matrixWorldInverse,
-      splat.matrixWorld,
-    );
-    material.defines.SORT32 = this.sort32;
-
     // Render
     const count = this.sort32 ? numSplats : numSplats / 2;
     const readbackBuffer = await this.ensureCapacity(count);
@@ -111,6 +104,17 @@ export class ReadbackSplatSorter implements SplatSorter {
 
     const renderState = this.saveRenderState(renderer);
 
+    const material = this.material;
+    splatData.setupMaterial(material);
+    material.uniforms.sortRadial.value = this.sortRadial;
+    material.uniforms.sortDepthBias.value = this.depthBias;
+    material.uniforms.sort360.value = this.sort360;
+    material.uniforms.splatModelViewMatrix.value.multiplyMatrices(
+      camera.matrixWorldInverse,
+      splat.matrixWorld,
+    );
+    material.defines.SORT32 = this.sort32;
+
     this.render(renderer, count, material);
     const promise = this.read(renderer, count, readback);
 
@@ -293,7 +297,7 @@ export class ReadbackSplatSorter implements SplatSorter {
     return Promise.all(promises).then(() => readback);
   }
 
-  private createMaterial(splatData: SplatData) {
+  private static createMaterial() {
     const shaders = getShaders();
     const material = new THREE.RawShaderMaterial({
       name: "SplatDistanceShader",
@@ -313,8 +317,6 @@ export class ReadbackSplatSorter implements SplatSorter {
       fragmentShader: shaders.splatDistanceFragment,
     });
 
-    splatData.setupMaterial(material);
-
     return material;
   }
 

src/encoding/ExtendedSplats.ts

@@ -209,9 +209,23 @@ export class ExtendedSplats implements IterableSplatData {
   }
 
   dispose(): void {
-    this.splatTexture1?.dispose();
-    this.splatTexture2?.dispose();
-    this.shTexture?.dispose();
+    if (this.splatTexture1) {
+      this.splatTexture1.dispose();
+      this.splatTexture1.source.data = null;
+    }
+    if (this.splatTexture2) {
+      this.splatTexture2.dispose();
+      this.splatTexture2.source.data = null;
+    }
+    if (this.shTexture) {
+      this.shTexture.dispose();
+      this.shTexture.source.data = null;
+    }
+    this.packedArray1 = EMPTY_UINT32_ARRAY;
+    this.packedArray2 = EMPTY_UINT32_ARRAY;
+    this.packedShArray = null;
+    this.numSplats = -1;
+    this.maxSplats = 0;
   }
 
   static encodingName = "extended";
@@ -447,3 +461,4 @@ export type EncodedExtendedSplats = {
 };
 
 const tempQuaternion = new THREE.Quaternion();
+const EMPTY_UINT32_ARRAY = new Uint32Array(0);

src/encoding/PackedSplats.ts

@@ -280,8 +280,18 @@ export class PackedSplats implements IterableSplatData {
   }
 
   dispose() {
-    this.texture?.dispose();
-    this.shTexture?.dispose();
+    if (this.texture) {
+      this.texture.dispose();
+      this.texture.source.data = null;
+    }
+    if (this.shTexture) {
+      this.shTexture.dispose();
+      this.shTexture.source.data = null;
+    }
+    this.packedArray = EMPTY_UINT32_ARRAY;
+    this.shArray = null;
+    this.numSplats = -1;
+    this.maxSplats = 0;
   }
 
   static encodingName = "packed";
@@ -536,3 +546,4 @@ export type EncodedPackedSplats = {
 };
 
 const tempQuaternion = new THREE.Quaternion();
+const EMPTY_UINT32_ARRAY = new Uint32Array(0);