More test coverage for argument origin tracking

iwoplaza · iwoplaza · commit 171af794f6c7 · 2025-11-06T15:28:30.000+01:00
diff --git a/packages/typegpu/src/core/function/fnCore.ts b/packages/typegpu/src/core/function/fnCore.ts
@@ -7,6 +7,7 @@ import {
   type Snippet,
 } from '../../data/snippet.ts';
 import {
+  isNaturallyEphemeral,
   isPtr,
   isWgslData,
   isWgslStruct,
@@ -196,39 +197,49 @@ export function createFnCore(
         // of the argument based on the argument's referentiality.
         // In other words, if we pass a reference to a function, it's typed as a pointer,
         // otherwise it's typed as a value.
-        const ref = isPtr(argType)
+        const origin = isPtr(argType)
           ? argType.addressSpace === 'storage'
             ? argType.access === 'read' ? 'readonly' : 'mutable'
             : argType.addressSpace
+          : isNaturallyEphemeral(argType)
+          ? 'runtime'
           : 'argument';
 
         switch (astParam?.type) {
           case FuncParameterType.identifier: {
             const rawName = astParam.name;
-            const snippet = snip(ctx.makeNameValid(rawName), argType, ref);
+            const snippet = snip(ctx.makeNameValid(rawName), argType, origin);
             args.push(snippet);
             if (snippet.value !== rawName) {
               argAliases.push([rawName, snippet]);
             }
             break;
           }
           case FuncParameterType.destructuredObject: {
-            args.push(snip(`_arg_${i}`, argType, ref));
-            argAliases.push(...astParam.props.map(({ name, alias }) =>
-              [
+            args.push(snip(`_arg_${i}`, argType, origin));
+            argAliases.push(...astParam.props.map(({ name, alias }) => {
+              // Undecorating, as the struct type can contain builtins
+              const destrType = undecorate(
+                (argTypes[i] as WgslStruct).propTypes[name],
+              );
+
+              const destrOrigin = isPtr(destrType)
+                ? destrType.addressSpace === 'storage'
+                  ? destrType.access === 'read' ? 'readonly' : 'mutable'
+                  : destrType.addressSpace
+                : isNaturallyEphemeral(destrType)
+                ? 'runtime'
+                : 'argument';
+
+              return [
                 alias,
-                snip(
-                  `_arg_${i}.${name}`,
-                  (argTypes[i] as WgslStruct)
-                    .propTypes[name],
-                  ref,
-                ),
-              ] as [string, Snippet]
-            ));
+                snip(`_arg_${i}.${name}`, destrType, destrOrigin),
+              ] as [string, Snippet];
+            }));
             break;
           }
           case undefined:
-            args.push(snip(`_arg_${i}`, argType, ref));
+            args.push(snip(`_arg_${i}`, argType, origin));
         }
       }
 
diff --git a/packages/typegpu/src/core/variable/tgpuVariable.ts b/packages/typegpu/src/core/variable/tgpuVariable.ts
@@ -1,5 +1,4 @@
 import type { AnyData } from '../../data/dataTypes.ts';
-import type { ref } from '../../data/ref.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
 import { isNaturallyEphemeral } from '../../data/wgslTypes.ts';
 import { IllegalVarAccessError } from '../../errors.ts';
@@ -26,7 +25,7 @@ export type VariableScope = 'private' | 'workgroup';
 export interface TgpuVar<
   TScope extends VariableScope = VariableScope,
   TDataType extends AnyData = AnyData,
-> extends TgpuNamable, ref<InferGPU<TDataType>> {
+> extends TgpuNamable {
   readonly [$gpuValueOf]: InferGPU<TDataType>;
   value: InferGPU<TDataType>;
   $: InferGPU<TDataType>;
diff --git a/packages/typegpu/src/data/ref.ts b/packages/typegpu/src/data/ref.ts
@@ -161,8 +161,10 @@ export function derefSnippet(snippet: Snippet): Snippet {
   invariant(isPtr(snippet.dataType), 'Only pointers can be dereferenced');
 
   const innerType = snippet.dataType.inner;
-  // Dereferencing a pointer does not return a copy of the value, it's still a reference.
-  const origin = isNaturallyEphemeral(innerType) ? 'runtime' : snippet.origin;
+  const origin =
+    isNaturallyEphemeral(innerType) && snippet.origin !== 'argument'
+      ? 'runtime'
+      : snippet.origin;
 
   if (snippet.value instanceof RefOperator) {
     return snip(stitch`${snippet.value.snippet}`, innerType, origin);
diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts
@@ -174,6 +174,8 @@ ${this.ctx.pre}}`;
       // Even types that aren't naturally referential (like vectors or structs) should
       // be treated as constant references when assigned to a const.
       varOrigin = 'constant-ref';
+    } else if (origin === 'argument' && !wgsl.isNaturallyEphemeral(dataType)) {
+      varOrigin = 'argument';
     } else if (!wgsl.isNaturallyEphemeral(dataType)) {
       varOrigin = isEphemeralOrigin(origin) ? 'function' : origin;
     } else if (origin === 'constant' && varType === 'const') {
@@ -291,11 +293,22 @@ ${this.ctx.pre}}`;
           );
         }
 
+        if (
+          rhsExpr.origin === 'argument' &&
+          !wgsl.isNaturallyEphemeral(rhsExpr.dataType)
+        ) {
+          throw new WgslTypeError(
+            `'${lhsStr} = ${rhsStr}' is invalid, because argument references cannot be assigned.\n-----\nTry '${lhsStr} = ${
+              this.ctx.resolve(rhsExpr.dataType).value
+            }(${rhsStr})' to copy the value instead.\n-----`,
+          );
+        }
+
         if (!isEphemeralSnippet(rhsExpr)) {
           throw new WgslTypeError(
             `'${lhsStr} = ${rhsStr}' is invalid, because references cannot be assigned.\n-----\nTry '${lhsStr} = ${
               this.ctx.resolve(rhsExpr.dataType).value
-            }(${rhsStr})' instead.\n-----`,
+            }(${rhsStr})' to copy the value instead.\n-----`,
           );
         }
       }
@@ -812,7 +825,10 @@ ${this.ctx.pre}else ${alternate}`;
 
       // Assigning a reference to a `const` variable means we store the pointer
       // of the rhs.
-      if (!isEphemeralSnippet(eq)) {
+      if (
+        !isEphemeralSnippet(eq) ||
+        (eq.origin === 'argument' && !wgsl.isNaturallyEphemeral(dataType))
+      ) {
         // Referential
         if (stmtType === NODE.let) {
           const rhsStr = this.ctx.resolve(eq.value).value;
diff --git a/packages/typegpu/tests/examples/individual/gravity.test.ts b/packages/typegpu/tests/examples/individual/gravity.test.ts
@@ -270,9 +270,9 @@ describe('gravity example', () => {
         var lightColor = vec3f(1, 0.8999999761581421, 0.8999999761581421);
         var textureColor = textureSample(celestialBodyTextures_9, sampler_10, input.uv, input.sphereTextureIndex).xyz;
         var ambient = ((textureColor * lightColor) * input.ambientLightFactor);
-        var normal = input.normals;
+        let normal = (&input.normals);
         var lightDirection = normalize((lightSource_11 - input.worldPosition));
-        let cosTheta = dot(normal, lightDirection);
+        let cosTheta = dot((*normal), lightDirection);
         var diffuse = ((textureColor * lightColor) * max(0f, cosTheta));
         var litColor = (ambient + diffuse);
         return vec4f(litColor.xyz, 1f);
diff --git a/packages/typegpu/tests/ref.test.ts b/packages/typegpu/tests/ref.test.ts
@@ -3,7 +3,7 @@ import { describe, expect } from 'vitest';
 import { it } from './utils/extendedIt.ts';
 import { asWgsl } from './utils/parseResolved.ts';
 
-describe('ref', () => {
+describe('d.ref', () => {
   it('fails when using a ref as an external', () => {
     const sup = d.ref(0);
 
diff --git a/packages/typegpu/tests/tgsl/argumentOrigin.test.ts b/packages/typegpu/tests/tgsl/argumentOrigin.test.ts
@@ -0,0 +1,148 @@
+import { describe, expect } from 'vitest';
+import * as d from '../../src/data/index.ts';
+import { it } from '../utils/extendedIt';
+import { asWgsl } from '../utils/parseResolved';
+
+describe('function argument origin tracking', () => {
+  it('should allow mutation of primitive arguments', () => {
+    const foo = (a: number) => {
+      'use gpu';
+      a += 1;
+    };
+
+    const main = () => {
+      'use gpu';
+      foo(1);
+    };
+
+    expect(asWgsl(main)).toMatchInlineSnapshot(`
+      "fn foo(a: i32) {
+        a += 1i;
+      }
+
+      fn main() {
+        foo(1i);
+      }"
+    `);
+  });
+
+  it('should allow mutation of destructured primitive arguments', () => {
+    const Foo = d.struct({ a: d.f32 });
+
+    const foo = ({ a }: { a: number }) => {
+      'use gpu';
+      a += 1;
+    };
+
+    const main = () => {
+      'use gpu';
+      foo(Foo({ a: 1 }));
+    };
+
+    expect(asWgsl(main)).toMatchInlineSnapshot(`
+      "struct Foo {
+        a: f32,
+      }
+
+      fn foo(_arg_0: Foo) {
+        _arg_0.a += 1f;
+      }
+
+      fn main() {
+        foo(Foo(1f));
+      }"
+    `);
+  });
+
+  it('should fail on mutation of non-primitive arguments', () => {
+    const foo = (a: d.v3f) => {
+      'use gpu';
+      a.x += 1;
+    };
+
+    const main = () => {
+      'use gpu';
+      foo(d.vec3f(1, 2, 3));
+    };
+
+    expect(() => asWgsl(main)).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:main
+      - fn*:main()
+      - fn*:foo(vec3f): 'a.x += 1f' is invalid, because non-pointer arguments cannot be mutated.]
+    `);
+  });
+
+  it('should fail on transitive mutation of non-primitive arguments', () => {
+    const foo = (a: d.v3f) => {
+      'use gpu';
+      const b = a;
+      b.x += 1;
+    };
+
+    const main = () => {
+      'use gpu';
+      foo(d.vec3f(1, 2, 3));
+    };
+
+    expect(() => asWgsl(main)).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:main
+      - fn*:main()
+      - fn*:foo(vec3f): '(*b).x += 1f' is invalid, because non-pointer arguments cannot be mutated.]
+    `);
+  });
+
+  it('should fail on create a let variable from an argument reference', () => {
+    const foo = (a: d.v3f) => {
+      'use gpu';
+      let b = a;
+      b = d.vec3f();
+      return b;
+    };
+
+    const main = () => {
+      'use gpu';
+      foo(d.vec3f(1, 2, 3));
+    };
+
+    expect(() => asWgsl(main)).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:main
+      - fn*:main()
+      - fn*:foo(vec3f): 'let b = a' is invalid, because references cannot be assigned to 'let' variable declarations.
+      -----
+      - Try 'let b = vec3f(a)' if you need to reassign 'b' later
+      - Try 'const b = a' if you won't reassign 'b' later.
+      -----]
+    `);
+  });
+
+  it('should fail on assigning an argument reference to a variable', () => {
+    const foo = (a: d.v3f) => {
+      'use gpu';
+      let b = d.vec3f();
+      b = a;
+      return b;
+    };
+
+    const main = () => {
+      'use gpu';
+      foo(d.vec3f(1, 2, 3));
+    };
+
+    expect(() => asWgsl(main)).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:main
+      - fn*:main()
+      - fn*:foo(vec3f): 'b = a' is invalid, because argument references cannot be assigned.
+      -----
+      Try 'b = vec3f(a)' to copy the value instead.
+      -----]
+    `);
+  });
+});
diff --git a/packages/typegpu/tests/tgslFn.test.ts b/packages/typegpu/tests/tgslFn.test.ts
@@ -164,8 +164,8 @@ describe('TGSL tgpu.fn function', () => {
       @vertex fn vertex_fn(input: vertex_fn_Input) -> vertex_fn_Output {
         let vi = f32(input.vi);
         let ii = f32(input.ii);
-        var color = input.color;
-        return vertex_fn_Output(vec4f(color.w, ii, vi, 1f), vec2f(color.w, vi));
+        let color = (&input.color);
+        return vertex_fn_Output(vec4f((*color).w, ii, vi, 1f), vec2f((*color).w, vi));
       }"
     `);
   });
@@ -391,9 +391,9 @@ describe('TGSL tgpu.fn function', () => {
       }
 
       @fragment fn fragmentFn(input: fragmentFn_Input) -> fragmentFn_Output {
-        var pos = input.pos;
+        let pos = (&input.pos);
         var sampleMask = 0;
-        if (((input.sampleMask > 0u) && (pos.x > 0f))) {
+        if (((input.sampleMask > 0u) && ((*pos).x > 0f))) {
           sampleMask = 1i;
         }
         return fragmentFn_Output(u32(sampleMask), 1f, vec4f());
