feat: Better snippet origin tracking

packages/typegpu/src/data/snippet.ts

@@ -1,7 +1,8 @@
-import { undecorate } from './dataTypes.ts';
-import type { AnyData, UnknownData } from './dataTypes.ts';
+import { stitch } from '../core/resolve/stitch.ts';
 import { DEV } from '../shared/env.ts';
-import { isNumericSchema } from './wgslTypes.ts';
+import type { AnyData, UnknownData } from './dataTypes.ts';
+import { undecorate } from './dataTypes.ts';
+import { isNaturallyEphemeral, isNumericSchema } from './wgslTypes.ts';
 
 export type Origin =
   | 'uniform'
@@ -29,12 +30,45 @@ export type Origin =
   | 'constant-tgpu-const-ref' /* turns into a `const` when assigned to a variable */
   | 'runtime-tgpu-const-ref' /* turns into a `let` when assigned to a variable */;
 
-export function isEphemeralOrigin(space: Origin) {
-  return space === 'runtime' || space === 'constant' || space === 'argument';
+export function isEphemeralSnippet(snippet: Snippet) {
+  if (snippet.origin === 'argument') {
+    // Arguments are considered ephemeral if their data type
+    // is naturally ephemeral.
+    // (primitives => true, non-primitives => false).
+    return isNaturallyEphemeral(snippet.dataType);
+  }
+  return snippet.origin === 'runtime' || snippet.origin === 'constant';
 }
 
-export function isEphemeralSnippet(snippet: Snippet) {
-  return isEphemeralOrigin(snippet.origin);
+/**
+ * Returns a reason if the snippet is immutable, or undefined if it is mutable.
+ */
+export function getImmutableReason(snippet: Snippet): string | undefined {
+  if (snippet.origin === 'argument') {
+    return stitch`Cannot mutate '${snippet}', arguments are immutable. You can copy them into a local variable first.`;
+  }
+
+  if (
+    snippet.origin === 'constant' ||
+    snippet.origin === 'constant-tgpu-const-ref' ||
+    snippet.origin === 'runtime-tgpu-const-ref'
+  ) {
+    return stitch`Cannot mutate '${snippet}', constants are immutable.`;
+  }
+
+  if (snippet.origin === 'uniform') {
+    return stitch`Cannot mutate '${snippet}', uniforms are immutable.`;
+  }
+
+  if (snippet.origin === 'handle') {
+    return stitch`Cannot mutate '${snippet}', handles are immutable.`;
+  }
+
+  if (snippet.origin === 'readonly') {
+    return stitch`Cannot mutate '${snippet}', it's bound as a readonly resource.`;
+  }
+
+  return undefined; // Mutable 🎉
 }
 
 export const originToPtrParams = {

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -11,10 +11,10 @@ import {
 } from '../data/dataTypes.ts';
 import { bool, i32, u32 } from '../data/numeric.ts';
 import {
-  isEphemeralOrigin,
   isEphemeralSnippet,
   isSnippet,
   type Origin,
+  type ResolvedSnippet,
   snip,
   type Snippet,
 } from '../data/snippet.ts';
@@ -224,41 +224,10 @@ ${this.ctx.pre}}`;
 
   public blockVariable(
     varType: 'var' | 'let' | 'const',
-    id: string,
-    dataType: wgsl.AnyWgslData | UnknownData,
-    origin: Origin,
-  ): Snippet {
-    const naturallyEphemeral = wgsl.isNaturallyEphemeral(dataType);
-
-    let varOrigin: Origin;
-    if (
-      origin === 'constant-tgpu-const-ref' ||
-      origin === 'runtime-tgpu-const-ref'
-    ) {
-      // Even types that aren't naturally referential (like vectors or structs) should
-      // be treated as constant references when assigned to a const.
-      varOrigin = origin;
-    } else if (origin === 'argument') {
-      if (naturallyEphemeral) {
-        varOrigin = 'runtime';
-      } else {
-        varOrigin = 'argument';
-      }
-    } else if (!naturallyEphemeral) {
-      varOrigin = isEphemeralOrigin(origin) ? 'this-function' : origin;
-    } else if (origin === 'constant' && varType === 'const') {
-      varOrigin = 'constant';
-    } else {
-      varOrigin = 'runtime';
-    }
-
-    const snippet = snip(
-      this.ctx.makeNameValid(id),
-      dataType,
-      /* origin */ varOrigin,
-    );
-    this.ctx.defineVariable(id, snippet);
-    return snippet;
+    lhs: ResolvedSnippet,
+    rhs: Snippet,
+  ): string {
+    return stitch`${this.ctx.pre}${varType} ${lhs.value} = ${rhs};`;
   }
 
   public identifier(id: string): Snippet {
@@ -916,10 +885,6 @@ ${this.ctx.pre}else ${alternate}`;
         );
       }
 
-      const ephemeral = isEphemeralSnippet(eq);
-      let dataType = eq.dataType as wgsl.AnyWgslData;
-      const naturallyEphemeral = wgsl.isNaturallyEphemeral(dataType);
-
       if (isLooseData(eq.dataType)) {
         throw new Error(
           `Cannot create variable '${rawId}' with loose data type.`,
@@ -947,10 +912,40 @@ ${this.ctx.pre}else ${alternate}`;
         };`;
       }
 
+      const ephemeral = isEphemeralSnippet(eq);
+      let dataType = eq.dataType as wgsl.AnyWgslData;
+      let varOrigin: Origin = 'runtime';
+      const naturallyEphemeral = wgsl.isNaturallyEphemeral(dataType);
+
+      if (eq.origin === 'argument') {
+        varOrigin = ephemeral ? 'runtime' : 'argument';
+
+        if (stmtType === NODE.let && !ephemeral) {
+          const rhsStr = this.ctx.resolve(eq.value).value;
+          const rhsTypeStr = this.ctx.resolve(unptr(eq.dataType)).value;
+
+          throw new WgslTypeError(
+            `'let ${rawId} = ${rhsStr}' is invalid, because references to arguments cannot be assigned to 'let' variable declarations.
+  -----
+  - Try 'let ${rawId} = ${rhsTypeStr}(${rhsStr})' if you need to reassign '${rawId}' later
+  - Try 'const ${rawId} = ${rhsStr}' if you won't reassign '${rawId}' later.
+  -----`,
+          );
+        }
+
+        if (stmtType === NODE.const) {
+          // Arguments cannot be mutated, so we 'let' them be (kill me)
+          varType = 'let';
+        }
+      } //
       // Assigning a reference to a `const` variable means we store the pointer
       // of the rhs.
-      if (!ephemeral) {
+      else if (!ephemeral) {
         // Referential
+
+        // It's a reference to something else, so we inherit their origin
+        varOrigin = eq.origin;
+
         if (stmtType === NODE.let) {
           const rhsStr = this.ctx.resolve(eq.value).value;
           const rhsTypeStr = this.ctx.resolve(unptr(eq.dataType)).value;
@@ -991,41 +986,41 @@ ${this.ctx.pre}else ${alternate}`;
       } else {
         // Non-referential
 
-        if (stmtType === NODE.const) {
-          if (eq.origin === 'argument') {
-            // Arguments cannot be mutated, so we 'let' them be (kill me)
-            varType = 'let';
-          } else if (naturallyEphemeral) {
+        if (naturallyEphemeral) {
+          // Primitives
+
+          if (stmtType === NODE.const) {
             varType = eq.origin === 'constant' ? 'const' : 'let';
+            varOrigin = eq.origin === 'constant' ? 'constant' : 'runtime';
           }
         } else {
-          // stmtType === NODE.let
-
-          if (eq.origin === 'argument') {
-            if (!naturallyEphemeral) {
-              const rhsStr = this.ctx.resolve(eq.value).value;
-              const rhsTypeStr = this.ctx.resolve(unptr(eq.dataType)).value;
+          // Non-primitives
 
-              throw new WgslTypeError(
-                `'let ${rawId} = ${rhsStr}' is invalid, because references to arguments cannot be assigned to 'let' variable declarations.
-  -----
-  - Try 'let ${rawId} = ${rhsTypeStr}(${rhsStr})' if you need to reassign '${rawId}' later
-  - Try 'const ${rawId} = ${rhsStr}' if you won't reassign '${rawId}' later.
-  -----`,
-              );
-            }
-          }
+          // We're creating a new variable with an ephemeral snippet,
+          // but referencing that variable down the line will have the
+          // origin of 'this-function'.
+          varOrigin = 'this-function';
         }
       }
 
-      const snippet = this.blockVariable(
+      if (
+        eq.origin === 'constant-tgpu-const-ref' ||
+        eq.origin === 'runtime-tgpu-const-ref'
+      ) {
+        // Even types that aren't naturally referential (like vectors or structs) should
+        // be treated as constant references when assigned to a const.
+        varOrigin = eq.origin;
+      }
+
+      const varName = this.ctx.makeNameValid(rawId);
+      const lhs = snip(varName, concretize(dataType), varOrigin);
+      this.ctx.defineVariable(rawId, lhs);
+
+      return this.blockVariable(
         varType,
-        rawId,
-        concretize(dataType),
-        eq.origin,
+        lhs,
+        tryConvertSnippet(eq, dataType, false),
       );
-      return stitch`${this.ctx.pre}${varType} ${snippet
-        .value as string} = ${tryConvertSnippet(eq, dataType, false)};`;
     }
 
     if (statement[0] === NODE.block) {

packages/typegpu/tests/bufferUsage.test.ts

@@ -88,6 +88,22 @@
         .as('uniform')
     ).toThrow();
   });
+
+  it('cannot be mutated', ({ root }) => {
+    const Boid = d.struct({
+      pos: d.vec3f,
+      vel: d.vec3u,
+    });
+
+    const boidUniform = root.createUniform(Boid);
+
+    const main = () => {
+      'use gpu';
+      boidUniform.$.pos.x += 1;
+    };
+
+    expect(() => tgpu.resolve([main])).toThrowErrorMatchingInlineSnapshot();
+  });
 });
 
 describe('TgpuBufferMutable', () => {

packages/typegpu/tests/constant.test.ts

@@ -121,49 +121,4 @@ describe('tgpu.const', () => {
       }"
     `);
   });
-
-  it('cannot be passed directly to shellless functions', () => {
-    const fn1 = (v: d.v3f) => {
-      'use gpu';
-      return v.x * v.y * v.z;
-    };
-
-    const foo = tgpu.const(d.vec3f, d.vec3f(1, 2, 3));
-    const fn2 = () => {
-      'use gpu';
-      return fn1(foo.$);
-    };
-
-    expect(() => tgpu.resolve([fn2])).toThrowErrorMatchingInlineSnapshot(`
-      [Error: Resolution of the following tree failed:
-      - <root>
-      - fn*:fn2
-      - fn*:fn2(): Cannot pass constant references as function arguments. Explicitly copy them by wrapping them in a schema: 'vec3f(...)']
-    `);
-  });
-
-  it('cannot be mutated', () => {
-    const boid = tgpu.const(Boid, {
-      pos: d.vec3f(1, 2, 3),
-      vel: d.vec3u(4, 5, 6),
-    });
-
-    const fn = () => {
-      'use gpu';
-      // @ts-expect-error: Cannot assign to read-only property
-      boid.$.pos = d.vec3f(0, 0, 0);
-    };
-
-    expect(() => tgpu.resolve([fn])).toThrowErrorMatchingInlineSnapshot(`
-      [Error: Resolution of the following tree failed:
-      - <root>
-      - fn*:fn
-      - fn*:fn(): 'boid.pos = vec3f()' is invalid, because boid.pos is a constant.]
-    `);
-
-    // Since we freeze the object, we cannot mutate when running the function in JS either
-    expect(() => fn()).toThrowErrorMatchingInlineSnapshot(
-      `[TypeError: Cannot assign to read only property 'pos' of object '#<Object>']`,
-    );
-  });
 });