Better handling of arguments

Author: iwoplaza

packages/typegpu-sdf/src/2d.ts

@@ -153,7 +153,7 @@ export const sdBezierApprox = tgpu.fn(
 });
 
 export const sdPie = tgpu.fn([vec2f, vec2f, f32], f32)((p, c, r) => {
-  const p_w = p;
+  const p_w = vec2f(p);
   p_w.x = abs(p.x);
   const l = length(p_w) - r;
   const m = length(p_w.sub(c.mul(clamp(dot(p_w, c), 0, r))));

packages/typegpu/src/data/ref.ts

@@ -1,5 +1,5 @@
 import { stitch } from '../core/resolve/stitch.ts';
-import { invariant } from '../errors.ts';
+import { invariant, WgslTypeError } from '../errors.ts';
 import { inCodegenMode } from '../execMode.ts';
 import { setName } from '../shared/meta.ts';
 import { $internal, $ownSnippet, $resolve } from '../shared/symbols.ts';
@@ -46,6 +46,12 @@ export interface ref<T> {
 
 export const ref: DualFn<<T>(value: T) => ref<T>> = (() => {
   const gpuImpl = (value: Snippet) => {
+    if (value.origin === 'argument') {
+      throw new WgslTypeError(
+        stitch`d.ref(${value}) is illegal, cannot take a reference of an argument. Copy the value locally first, and take a reference of the copy.`,
+      );
+    }
+
     /**
      * Pointer type only exists if the ref was created from a reference (buttery-butter).
      *
@@ -161,10 +167,7 @@ export function derefSnippet(snippet: Snippet): Snippet {
   invariant(isPtr(snippet.dataType), 'Only pointers can be dereferenced');
 
   const innerType = snippet.dataType.inner;
-  const origin =
-    isNaturallyEphemeral(innerType) && snippet.origin !== 'argument'
-      ? 'runtime'
-      : snippet.origin;
+  const origin = isNaturallyEphemeral(innerType) ? 'runtime' : snippet.origin;
 
   if (snippet.value instanceof RefOperator) {
     return snip(stitch`${snippet.value.snippet}`, innerType, origin);

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -717,6 +717,26 @@ ${this.ctx.pre}}`;
           );
         }
 
+        // Arguments cannot be returned from functions without copying. A simple example why is:
+        // const identity = (x) => {
+        //   'use gpu';
+        //   return x;
+        // };
+        //
+        // const foo = (arg: d.v3f) => {
+        //   'use gpu';
+        //   const marg = identity(arg);
+        //   marg.x = 1; // 'marg's origin would be 'runtime', so we wouldn't be able to track this misuse.
+        // };
+        if (
+          returnSnippet.origin === 'argument' &&
+          !isEphemeralSnippet(returnSnippet)
+        ) {
+          throw new WgslTypeError(
+            stitch`Cannot return references to arguments, returning '${returnSnippet}'. Copy the argument before returning it.`,
+          );
+        }
+
         if (
           !expectedReturnType &&
           !isEphemeralSnippet(returnSnippet) &&
@@ -794,6 +814,10 @@ ${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.`,
@@ -821,14 +845,9 @@ ${this.ctx.pre}else ${alternate}`;
         };`;
       }
 
-      let dataType = eq.dataType as wgsl.AnyWgslData;
-
       // Assigning a reference to a `const` variable means we store the pointer
       // of the rhs.
-      if (
-        !isEphemeralSnippet(eq) ||
-        (eq.origin === 'argument' && !wgsl.isNaturallyEphemeral(dataType))
-      ) {
+      if (!ephemeral) {
         // Referential
         if (stmtType === NODE.let) {
           const rhsStr = this.ctx.resolve(eq.value).value;
@@ -866,11 +885,32 @@ ${this.ctx.pre}else ${alternate}`;
         }
       } else {
         // Non-referential
-        if (
-          stmtType === NODE.const &&
-          wgsl.isNaturallyEphemeral(dataType)
-        ) {
-          varType = eq.origin === 'constant' ? 'const' : 'let';
+
+        if (eq.origin === 'argument') {
+          if (stmtType === NODE.let) {
+            const rhsStr = this.ctx.resolve(eq.value).value;
+            const rhsTypeStr =
+              this.ctx.resolve(toStorable(eq.dataType as wgsl.StorableData))
+                .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.
+-----`,
+            );
+          }
+          varType = 'let';
+        }
+
+        if (stmtType === NODE.const) {
+          if (eq.origin === 'argument') {
+            // Arguments cannot be mutated, so we 'let' them be (kill me)
+            varType = 'let';
+          } else if (naturallyEphemeral) {
+            varType = eq.origin === 'constant' ? 'const' : 'let';
+          }
         }
       }
 

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);
-        let 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);