fix(@typegpu/gl): release unused uniforms and destroyed textures

packages/typegpu-gl/src/tgpuRootWebGL.ts

@@ -179,7 +179,6 @@ class TgpuWebGLRenderPipelineImpl implements TgpuWebGLRenderPipeline {
     gl: WebGL2RenderingContext,
     program: WebGLProgram,
     crossShaderStageState: CrossShaderStageState,
-    uniforms: readonly WebGLUniform[],
     offscreen: OffscreenCanvas,
   ) {
     this.#gl = gl;
@@ -189,13 +188,10 @@ class TgpuWebGLRenderPipelineImpl implements TgpuWebGLRenderPipeline {
     if (!vao) throw new Error('Failed to create VAO');
     this.#vao = vao;
 
-    // Query uniform locations once; skip uniforms that weren't actually used by the shaders.
+    // Discover bindings from shader references, without retaining a root-wide uniform history.
     const bindings: UniformBinding[] = [];
-    for (const uniform of uniforms) {
-      const name = crossShaderStageState.globalIdentifierMap.get(uniform);
-      if (!name) {
-        continue; // Not used in the shader
-      }
+    for (const [uniform, name] of crossShaderStageState.globalIdentifierMap) {
+      if (!(uniform instanceof WebGLUniformImpl)) continue;
 
       const location = gl.getUniformLocation(program, name);
       if (location === null) {
@@ -406,9 +402,8 @@ class WebGLUniformImpl<TData extends d.AnyWgslData> implements WebGLUniform<TDat
 export class TgpuRootWebGL {
   #gl: WebGL2RenderingContext;
   #offscreen: OffscreenCanvas;
-  #uniforms: WebGLUniformImpl<d.AnyWgslData>[] = [];
   #buffers: WebGLBuffer[] = [];
-  #textures: WebGLTextureImpl[] = [];
+  readonly #textures = new Set<WebGLTextureImpl>();
   #samplers: WebGLSamplerImpl[] = [];
   #programCache = new Map<string, WebGLProgram>();
 
@@ -425,9 +420,7 @@ export class TgpuRootWebGL {
     typeSchema: TData,
     initial?: BufferInitialData<TData>,
   ): WebGLUniform<TData> {
-    const uniform = new WebGLUniformImpl(typeSchema, initial);
-    this.#uniforms.push(uniform as unknown as WebGLUniformImpl<d.AnyWgslData>);
-    return uniform;
+    return new WebGLUniformImpl(typeSchema, initial);
   }
 
   createMutable(): never {
@@ -463,8 +456,12 @@ export class TgpuRootWebGL {
   }
 
   createTexture<TProps extends TextureProps>(props: TProps): TgpuTexture<TProps> {
-    const texture = new WebGLTextureImpl(this.#gl, props);
-    this.#textures.push(texture);
+    // A surviving texture must not keep the registry and its other textures alive.
+    const textures = new WeakRef(this.#textures);
+    const texture = new WebGLTextureImpl(this.#gl, props, (destroyed) => {
+      textures.deref()?.delete(destroyed);
+    });
+    this.#textures.add(texture);
     return asTgpuTexture<TProps>(texture);
   }
 
@@ -529,7 +526,6 @@ export class TgpuRootWebGL {
       this.#gl,
       program,
       crossShaderStageState,
-      this.#uniforms.slice() as Array<WebGLUniform>,
       this.#offscreen,
     );
   }
@@ -563,14 +559,10 @@ export class TgpuRootWebGL {
       this.#gl.deleteBuffer(buf);
     }
     this.#buffers = [];
-    for (const uniform of this.#uniforms) {
-      uniform.destroy();
-    }
-    this.#uniforms = [];
     for (const texture of this.#textures) {
       texture.destroy();
     }
-    this.#textures = [];
+    this.#textures.clear();
     for (const sampler of this.#samplers) {
       this.#gl.deleteSampler(sampler.raw);
     }

packages/typegpu-gl/src/webglTexture.ts

@@ -227,8 +227,13 @@ export class WebGLTextureImpl {
   needsYFlipWhenSampling = false;
   destroyed = false;
   readonly #framebuffers = new Set<WebGLFramebuffer>();
+  #onDestroy: ((texture: WebGLTextureImpl) => void) | undefined;
 
-  constructor(gl: WebGL2RenderingContext, props: TextureProps) {
+  constructor(
+    gl: WebGL2RenderingContext,
+    props: TextureProps,
+    onDestroy?: (texture: WebGLTextureImpl) => void,
+  ) {
     const depth = props.size[2] ?? 1;
     if ((props.dimension ?? '2d') !== '2d' || depth !== 1) {
       throw new Error('WebGL fallback currently supports only 2D textures with one layer.');
@@ -239,6 +244,7 @@ export class WebGLTextureImpl {
 
     this.gl = gl;
     this.props = props;
+    this.#onDestroy = onDestroy;
     this.width = props.size[0] ?? 1;
     this.height = props.size[1] ?? 1;
     this.format = getFormat(gl, props.format);
@@ -380,6 +386,8 @@ export class WebGLTextureImpl {
     }
     this.#framebuffers.clear();
     this.gl.deleteTexture(this.raw);
+    this.#onDestroy?.(this);
+    this.#onDestroy = undefined;
   }
 }
 

packages/typegpu-gl/tests/garbageCollection.test.ts

@@ -0,0 +1,139 @@
+import * as v8 from 'node:v8';
+import * as vm from 'node:vm';
+import { describe, expect, vi } from 'vitest';
+import { d } from 'typegpu';
+import { initWithGL } from '../src/index.ts';
+import { TgpuRootWebGL, type TgpuWebGLRenderPipeline } from '../src/tgpuRootWebGL.ts';
+import { it } from './utils/extendedTest.ts';
+
+v8.setFlagsFromString('--expose-gc');
+const gc = vm.runInNewContext('gc') as () => void;
+
+async function collect(refs: WeakRef<object>[]): Promise<number> {
+  let alive = refs;
+  for (let i = 0; i < 20 && alive.length > 0; i++) {
+    // WeakRef targets stay alive until the end of the current job.
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    gc();
+    alive = alive.filter((ref) => ref.deref() !== undefined);
+  }
+  return alive.length;
+}
+
+describe('TgpuRootWebGL - garbage collection', () => {
+  it('collects unused uniforms and their buffers while their root stays alive', async ({ gl }) => {
+    const root = new TgpuRootWebGL(gl);
+    const refs = Array.from({ length: 100 }, () => {
+      const uniform = root.createUniform(d.vec4f);
+      return [new WeakRef(uniform), new WeakRef(uniform.buffer)];
+    }).flat();
+
+    expect(await collect(refs)).toBe(0);
+    root.destroy();
+  });
+
+  it('keeps uniforms alive only while a pipeline uses them', async ({ gl }) => {
+    const root = new TgpuRootWebGL(gl);
+    function createPipeline(): {
+      refs: WeakRef<object>[];
+      pipeline: TgpuWebGLRenderPipeline | undefined;
+    } {
+      const uniform = root.createUniform(d.f32, 2);
+      return {
+        refs: [new WeakRef(uniform), new WeakRef(uniform.buffer)],
+        pipeline: root.createRenderPipeline({
+          vertex: () => {
+            'use gpu';
+            return { $position: d.vec4f(0, 0, 0, 1) };
+          },
+          fragment: () => {
+            'use gpu';
+            return d.vec4f(uniform.$);
+          },
+        }),
+      };
+    }
+    const resources = createPipeline();
+
+    expect(await collect(resources.refs)).toBe(2);
+    resources.pipeline?.draw(3);
+    expect(gl.uniform1f).toHaveBeenCalledExactlyOnceWith(expect.anything(), 2);
+
+    resources.pipeline = undefined;
+    expect(await collect(resources.refs)).toBe(0);
+    root.destroy();
+  });
+
+  it('collects destroyed textures while their root stays alive', async ({ gl }) => {
+    const root = initWithGL({ gl });
+    const refs = Array.from({ length: 100 }, () => {
+      const texture = root.createTexture({ size: [1, 1], format: 'rgba8unorm' });
+      texture.destroy();
+      return new WeakRef(texture);
+    });
+
+    expect(await collect(refs)).toBe(0);
+    root.destroy();
+    expect(gl.deleteTexture).toHaveBeenCalledTimes(100);
+  });
+
+  it('collects optimized-out uniforms while their pipeline stays alive', async ({ gl }) => {
+    const root = new TgpuRootWebGL(gl);
+    vi.mocked(gl.getUniformLocation).mockReturnValue(null);
+    function createPipeline() {
+      const uniform = root.createUniform(d.f32, 2);
+      return {
+        refs: [new WeakRef(uniform), new WeakRef(uniform.buffer)],
+        pipeline: root.createRenderPipeline({
+          vertex: () => {
+            'use gpu';
+            return { $position: d.vec4f(0, 0, 0, 1) };
+          },
+          fragment: () => {
+            'use gpu';
+            // oxlint-disable-next-line oxc/erasing-op -- Exercise native optimization of a referenced uniform.
+            return d.vec4f(uniform.$ * 0);
+          },
+        }),
+      };
+    }
+    const { refs, pipeline } = createPipeline();
+
+    expect(gl.getUniformLocation).toHaveBeenCalledTimes(1);
+    expect(await collect(refs)).toBe(0);
+    pipeline.draw(3);
+    expect(gl.uniform1f).not.toHaveBeenCalled();
+    expect(gl.drawArrays).toHaveBeenCalledTimes(1);
+    root.destroy();
+  });
+
+  it('does not let a retained texture keep its root or sibling textures alive', async ({ gl }) => {
+    function createTextures() {
+      const root = initWithGL({ gl });
+      const kept = root.createTexture({ size: [1, 1], format: 'rgba8unorm' });
+      const sibling = root.createTexture({ size: [1, 1], format: 'rgba8unorm' });
+      return { kept, refs: [new WeakRef(root), new WeakRef(sibling)] };
+    }
+    const { kept, refs } = createTextures();
+
+    expect(await collect(refs)).toBe(0);
+    kept.destroy();
+    expect(gl.deleteTexture).toHaveBeenCalledTimes(1);
+  });
+
+  it('releases cleanup references even when a destroyed texture is retained', async ({ gl }) => {
+    function createTextures() {
+      const root = initWithGL({ gl });
+      const kept = root.createTexture({ size: [1, 1], format: 'rgba8unorm' });
+      const sibling = root.createTexture({ size: [1, 1], format: 'rgba8unorm' });
+      const refs = [new WeakRef(root), new WeakRef(sibling)];
+      kept.destroy();
+      return { kept, refs };
+    }
+    const { kept, refs } = createTextures();
+
+    expect(await collect(refs)).toBe(0);
+    kept.destroy();
+    expect(gl.deleteTexture).toHaveBeenCalledTimes(1);
+  });
+});

packages/typegpu-gl/tests/uniforms.test.ts

@@ -1,4 +1,4 @@
-import { describe, beforeEach, expect } from 'vitest';
+import { describe, expect, vi } from 'vitest';
 import { tgpu, d } from 'typegpu';
 import { dualGlOptions, glOptions, initWithGL } from '@typegpu/gl';
 import { it } from './utils/extendedTest.ts';
@@ -35,6 +35,85 @@ describe('TgpuRootWebGL - createUniform', () => {
   });
 });
 
+describe('TgpuRootWebGL - uniform bindings', () => {
+  it('binds referenced uniforms from both stages and uploads subsequent writes', ({ gl }) => {
+    const root = initWithGL({ gl });
+    root.createUniform(d.f32, 99).$name('unused');
+    const shared = root.createUniform(d.f32, 1).$name('factor');
+    const offset = root.createUniform(d.f32, 0.5).$name('offset');
+    const color = root.createUniform(d.vec4f, d.vec4f(1, 0, 0, 1)).$name('color');
+    const optimizedOut = root.createUniform(d.f32, 2).$name('optimizedOut');
+    vi.mocked(gl.getUniformLocation).mockImplementation((program, name) =>
+      name === 'optimizedOut' ? null : ({ program, name } as unknown as WebGLUniformLocation),
+    );
+    const descriptor = {
+      vertex: () => {
+        'use gpu';
+        return { $position: d.vec4f(shared.$ + offset.$, 0, 0, 1) };
+      },
+      fragment: () => {
+        'use gpu';
+        return color.$ * shared.$ * optimizedOut.$;
+      },
+    };
+    // Each upload must target a location queried on the program bound for that draw. This holds
+    // whether or not pipelines with identical GLSL share one linked program.
+    function uploadedValues(): Record<string, unknown> {
+      const drawnProgram = vi.mocked(gl.useProgram).mock.lastCall?.[0];
+      expect(drawnProgram).toBeDefined();
+      const uploads = [
+        ...vi.mocked(gl.uniform1f).mock.calls,
+        ...vi.mocked(gl.uniform4fv).mock.calls,
+      ].map(([location, value]) => {
+        const { program, name } = location as unknown as { program: WebGLProgram; name: string };
+        expect(program).toBe(drawnProgram);
+        return [name, value] as const;
+      });
+      expect(uploads).toHaveLength(3);
+      return Object.fromEntries(uploads);
+    }
+
+    const pipeline = root.createRenderPipeline(descriptor);
+    expect(
+      vi
+        .mocked(gl.getUniformLocation)
+        .mock.calls.map(([, name]) => name)
+        .toSorted(),
+    ).toEqual(['color', 'factor', 'offset', 'optimizedOut']);
+
+    shared.write(3);
+    color.write(d.vec4f(0, 1, 0, 1));
+    pipeline.draw(3);
+    expect(uploadedValues()).toEqual({
+      factor: 3,
+      offset: 0.5,
+      color: new Float32Array([0, 1, 0, 1]),
+    });
+
+    // A uniform can be reused by a later pipeline with a different program, too.
+    vi.mocked(gl.uniform1f).mockClear();
+    vi.mocked(gl.uniform4fv).mockClear();
+    offset.write(0.75);
+    root
+      .createRenderPipeline({
+        vertex: descriptor.vertex,
+        fragment: () => {
+          'use gpu';
+          return color.$ * shared.$;
+        },
+      })
+      .draw(3);
+    expect(vi.mocked(gl.useProgram).mock.lastCall?.[0]).not.toBe(
+      vi.mocked(gl.useProgram).mock.calls[0]?.[0],
+    );
+    expect(uploadedValues()).toEqual({
+      factor: 3,
+      offset: 0.75,
+      color: new Float32Array([0, 1, 0, 1]),
+    });
+  });
+});
+
 describe('GlslGenerator - uniform resolution', () => {
   it('emits a uniform declaration and references the name in shader body', ({ gl }) => {
     const root = initWithGL({ gl });

packages/typegpu-gl/tests/webglFallback.test.ts

@@ -141,6 +141,44 @@ describe('TgpuRootWebGL - createRenderPipeline', () => {
 });
 
 describe('TgpuRootWebGL - textures', () => {
+  it('destroys remaining textures and their framebuffers exactly once', ({ gl }) => {
+    const root = initWithGL({ gl });
+    const textures = Array.from({ length: 3 }, () =>
+      root.createTexture({ size: [1, 1], format: 'rgba8unorm' }).$usage('render'),
+    );
+    for (const texture of textures) {
+      texture.createView('render');
+    }
+
+    textures[0]?.destroy();
+    root.destroy();
+    for (const texture of textures) {
+      expect(texture.destroyed).toBe(true);
+    }
+
+    const nativeTextures = vi.mocked(gl.createTexture).mock.results.map((result) => result.value);
+    const framebuffers = vi.mocked(gl.createFramebuffer).mock.results.map((result) => result.value);
+    expect(gl.deleteTexture).toHaveBeenCalledTimes(nativeTextures.length);
+    expect(gl.deleteFramebuffer).toHaveBeenCalledTimes(framebuffers.length);
+    for (const raw of nativeTextures) {
+      expect(
+        vi.mocked(gl.deleteTexture).mock.calls.filter(([deleted]) => deleted === raw),
+      ).toHaveLength(1);
+    }
+    for (const framebuffer of framebuffers) {
+      expect(
+        vi.mocked(gl.deleteFramebuffer).mock.calls.filter(([deleted]) => deleted === framebuffer),
+      ).toHaveLength(1);
+    }
+
+    root.destroy();
+    for (const texture of textures) {
+      texture.destroy();
+    }
+    expect(gl.deleteTexture).toHaveBeenCalledTimes(nativeTextures.length);
+    expect(gl.deleteFramebuffer).toHaveBeenCalledTimes(framebuffers.length);
+  });
+
   it('uploads CPU data to a 2D texture', ({ gl }) => {
     const root = initWithGL({ gl });
     const texture = root.createTexture({ size: [2, 2], format: 'rgba8unorm' });