feat(@typegpu/gl): Implicit pointer definitions in GLSL generator

packages/typegpu-gl/src/glslGenerator.ts

@@ -1,5 +1,5 @@
 import { NodeTypeCatalog as NODE } from 'tinyest';
-import type { Return } from 'tinyest';
+import type { Expression, Return } from 'tinyest';
 import { tgpu, d, type ShaderStage } from 'typegpu';
 import {
   abstractInt,
@@ -14,6 +14,7 @@ import type {
   FunctionDefinitionOptions,
   ConstantDefinitionOptions,
   VariableDefinitionOptions,
+  Origin,
   Snippet,
   ResolvedSnippet,
 } from 'typegpu/~internal';
@@ -119,6 +120,12 @@ interface EntryFnState {
   fragColorName?: string;
 }
 
+/**
+ * Origins of values that cannot be mutated for the whole duration of a shader's
+ * execution. Taking a reference to them is equivalent to copying them.
+ */
+const immutableOrigins: readonly Origin[] = ['uniform', 'readonly', 'handle'];
+
 function undecorateDataType(t: d.BaseData): d.BaseData {
   return d.isDecorated(t) ? t.inner : t;
 }
@@ -208,7 +215,7 @@ export class GlslGenerator extends WgslGenerator {
     this.#vertexOutPropToVarMap = {};
   }
 
-  public initGenerator(ctx: ResolutionCtx): void {
+  override initGenerator(ctx: ResolutionCtx): void {
     super.initGenerator(ctx);
     ctxToCrossShaderStageStateMap.set(ctx, this.#crossShaderStageState);
 
@@ -423,6 +430,82 @@ export class GlslGenerator extends WgslGenerator {
     return `${this.ctx.pre}${glslTypeName} ${name}${resolveArraySizeSuffix(this.ctx, dataType)} = ${rhsStr};`;
   }
 
+  /**
+   * GLSL has no pointers, so `const x = <alias>;` cannot be turned into an implicit
+   * pointer definition like it is in WGSL. Instead:
+   * - if the aliased memory is immutable for the whole shader run (uniforms, ...), we
+   *   copy the value, which is indistinguishable from referencing it,
+   * - otherwise `x` becomes an alias, meaning every use of it is replaced with the
+   *   expression it points to. Index expressions are hoisted into variables first, so
+   *   that they're evaluated exactly once, at the point of the declaration.
+   */
+  protected override _aliasConstStatement(rawId: string, eqNode: Expression, eq: Snippet): string {
+    if (immutableOrigins.includes(eq.origin)) {
+      const dataType = eq.dataType as d.BaseData;
+      const name = this.ctx.makeUniqueIdentifier(rawId, 'block');
+      this.ctx.defineVariable(rawId, snip(name, dataType, 'runtime-immutable-def', false));
+      return this._emitVarDecl('let', name, dataType, this.ctx.resolveSnippet(eq).value);
+    }
+
+    // The aliased memory can change over time, so copying would alter the semantics.
+    const hoisted: string[] = [];
+    const aliased = this._expression(this.#hoistIndexAccesses(eqNode, hoisted));
+
+    this.ctx.defineVariable(
+      rawId,
+      snip(
+        this.ctx.resolveSnippet(aliased).value,
+        aliased.dataType as d.BaseData,
+        aliased.origin,
+        false,
+      ),
+    );
+
+    return hoisted.join('\n');
+  }
+
+  /**
+   * Replaces every index expression in `node` that could change value over time with a
+   * reference to a freshly declared variable, whose declaration is appended to `out`.
+   *
+   * @example
+   * ```
+   * arr[foo()].prop[idx]  =>  arr[item].prop[item_1]
+   * // out: ['int item = foo();', 'int item_1 = idx;']
+   * ```
+   */
+  #hoistIndexAccesses(node: Expression, out: string[]): Expression {
+    if (typeof node !== 'object') {
+      return node;
+    }
+
+    if (node[0] === NODE.memberAccess) {
+      return [NODE.memberAccess, this.#hoistIndexAccesses(node[1], out), node[2]];
+    }
+
+    if (node[0] === NODE.indexAccess) {
+      const target = this.#hoistIndexAccesses(node[1], out);
+      const index = this._expression(node[2]);
+
+      if (
+        !index.possibleSideEffects &&
+        (index.origin === 'constant' || index.origin === 'constant-immutable-def')
+      ) {
+        // Known at comptime, so it cannot change between now and the uses of the alias.
+        // It also cannot have side-effects, so it can just be copied in many places.
+        return [NODE.indexAccess, target, node[2]];
+      }
+
+      const resolved = this.ctx.resolveSnippet(index);
+      const name = this.ctx.makeUniqueIdentifier('idx', 'block');
+      out.push(this._emitVarDecl('let', name, resolved.dataType, resolved.value));
+      this.ctx.defineVariable(name, snip(name, resolved.dataType, 'runtime-immutable-def', false));
+      return [NODE.indexAccess, target, name];
+    }
+
+    return node;
+  }
+
   override _return(statement: Return): string {
     const exprNode = statement[1];
 

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

@@ -0,0 +1,242 @@
+// oxlint-disable typescript/no-unnecessary-type-assertion
+import { describe, expect } from 'vitest';
+import { d, tgpu } from 'typegpu';
+import { glOptions } from '@typegpu/gl';
+import { it } from './utils/extendedTest.ts';
+import { initWithGL } from '../src/initWithGL.ts';
+
+const Boid = d.struct({
+  pos: d.vec3f,
+  vel: d.vec3f,
+});
+
+describe('implicit pointers in GLSL', () => {
+  it('copies references to immutable memory', ({ gl }) => {
+    const root = initWithGL({ gl });
+    const boid = root.createUniform(Boid);
+
+    function foo() {
+      'use gpu';
+      const boidPos = boid.$.pos;
+      return d.vec3f(boidPos);
+    }
+
+    expect(tgpu.resolve([foo], glOptions())).toMatchInlineSnapshot(`
+      "struct Boid {
+        vec3 pos;
+        vec3 vel;
+      };
+
+      uniform Boid boid;
+
+      vec3 foo() {
+        vec3 boidPos = boid.pos;
+        return boidPos;
+      }"
+    `);
+  });
+
+  it('aliases references to mutable memory', () => {
+    const boids = tgpu.privateVar(d.arrayOf(Boid, 16));
+
+    function bar() {
+      'use gpu';
+      const boid = boids.$[0] as d.Infer<typeof Boid>;
+      boid.pos.x += 1;
+      boid.vel = d.vec3f();
+    }
+
+    expect(tgpu.resolve([bar], glOptions())).toMatchInlineSnapshot(`
+      "struct Boid {
+        vec3 pos;
+        vec3 vel;
+      };
+
+      Boid boids[16];
+
+      void bar() {
+        boids[0].pos.x += 1.0;
+        boids[0].vel = vec3(0);
+      }"
+    `);
+  });
+
+  it('stores runtime index expressions in variables', () => {
+    const boids = tgpu.privateVar(d.arrayOf(Boid, 16));
+
+    function firstIndex() {
+      'use gpu';
+      return 0;
+    }
+
+    function bar(index: number) {
+      'use gpu';
+      const boid = boids.$[firstIndex() + index] as d.Infer<typeof Boid>;
+      boid.pos.x += 1;
+      boid.pos.y += 1;
+    }
+
+    function main() {
+      'use gpu';
+      bar(1);
+    }
+
+    expect(tgpu.resolve([main], glOptions())).toMatchInlineSnapshot(`
+      "int firstIndex() {
+        return 0;
+      }
+
+      struct Boid {
+        vec3 pos;
+        vec3 vel;
+      };
+
+      Boid boids[16];
+
+      void bar(int index) {
+        int idx = (firstIndex() + index);
+        boids[idx].pos.x += 1.0;
+        boids[idx].pos.y += 1.0;
+      }
+
+      void main() {
+        bar(1);
+      }"
+    `);
+  });
+
+  it('aliases nested member and index accesses', () => {
+    const Cluster = d.struct({
+      boids: d.arrayOf(Boid, 4),
+    });
+    const clusters = tgpu.privateVar(d.arrayOf(Cluster, 2));
+
+    function bar(index: number) {
+      'use gpu';
+      const cluster = clusters.$[index];
+      const pos = cluster!.boids[1]!.pos;
+      pos.x = 1;
+    }
+
+    function main() {
+      'use gpu';
+      bar(1);
+    }
+
+    expect(tgpu.resolve([main], glOptions())).toMatchInlineSnapshot(`
+      "struct Boid {
+        vec3 pos;
+        vec3 vel;
+      };
+
+      struct Cluster {
+        Boid boids[4];
+      };
+
+      Cluster clusters[2];
+
+      void bar(int index) {
+        int idx = index;
+        clusters[idx].boids[1].pos.x = 1.0;
+      }
+
+      void main() {
+        bar(1);
+      }"
+    `);
+  });
+
+  it('aliases an alias', () => {
+    const boids = tgpu.privateVar(d.arrayOf(Boid, 16));
+
+    function bar(index: number) {
+      'use gpu';
+      const boid = boids.$[index]!;
+      const pos = boid.pos;
+      pos.x = 1;
+    }
+
+    function main() {
+      'use gpu';
+      bar(1);
+    }
+
+    expect(tgpu.resolve([main], glOptions())).toMatchInlineSnapshot(`
+      "struct Boid {
+        vec3 pos;
+        vec3 vel;
+      };
+
+      Boid boids[16];
+
+      void bar(int index) {
+        int idx = index;
+        boids[idx].pos.x = 1.0;
+      }
+
+      void main() {
+        bar(1);
+      }"
+    `);
+  });
+
+  it('aliases a local variable', () => {
+    function bar() {
+      'use gpu';
+      const boid = Boid();
+      const pos = boid.pos;
+      pos.x = 1;
+      return boid.pos.x;
+    }
+
+    expect(tgpu.resolve([bar], glOptions())).toMatchInlineSnapshot(`
+      "struct Boid {
+        vec3 pos;
+        vec3 vel;
+      };
+
+      float bar() {
+        Boid boid = Boid();
+        boid.pos.x = 1.0;
+        return boid.pos.x;
+      }"
+    `);
+  });
+
+  it('hoists multiple index accessed', () => {
+    const boids = tgpu.privateVar(d.arrayOf(d.arrayOf(Boid, 16), 16));
+
+    function bar(index: number) {
+      'use gpu';
+      const idx2 = 4;
+      const boid = boids.$[index]![idx2]!;
+      const pos = boid.pos;
+      pos.x = 1;
+    }
+
+    function main() {
+      'use gpu';
+      bar(1);
+    }
+
+    expect(tgpu.resolve([main], glOptions())).toMatchInlineSnapshot(`
+      "struct Boid {
+        vec3 pos;
+        vec3 vel;
+      };
+
+      Boid boids[16][16];
+
+      void bar(int index) {
+        int idx2 = 4;
+        int idx = index;
+        int idx_1 = idx2;
+        boids[idx][idx_1].pos.x = 1.0;
+      }
+
+      void main() {
+        bar(1);
+      }"
+    `);
+  });
+});