refactor: Better origin tracking (#2490)

Author: iwoplaza

apps/typegpu-docs/src/content/docs/advanced/shader-generation.mdx

@@ -197,10 +197,10 @@ There are essentially three types of origins:
   performing operations that produce new values. Examples include creating a new `Boid` instance or calculating a new position based on an existing one. These
   include `'runtime'` and `'constant'`.
 - **Referential Origins**: These origins represent references to existing values. They are typically used for accessing or modifying existing data. Examples
-  include accessing the position of an existing `Boid` instance or modifying the position of an existing `Boid` instance. These include `'uniform'`, `'mutable'`, `'readonly'`, `'workgroup'`, `'private'`, `'function'`, `'handle'`, `'constant-ref'` and `'this-function'`.
+  include accessing the position of an existing `Boid` instance or modifying the position of an existing `Boid` instance:
     - `'uniform'`, `'mutable'`, `'readonly'`, `'workgroup'`, `'private'`, `'function'`, `'handle'` all reflect address spaces that values can belong to, and
       we use them to determine what kind of pointer type they are.
-    - `'constant-tgpu-const-ref'` is a reference to a value stored in a `tgpu.const`. They're different from `constant`s, as we know that even if they're referential (non-primitive), the developer cannot mutate them.
-    - `'runtime-tgpu-const-ref'` is a reference to a value stored in a `tgpu.const` but accessed in a way that isn’t constant at shader‑creation time (for example, indexed using a runtime value).
-    - `'this-function'` lets us track whether values originates from the function we're currently generating, or the function that called us.
+    - `'constant-immutable-def'` is a reference to a value stored in a `tgpu.const`. They're different from `constant`s, as we know that even if they're referential (non-primitive), the developer cannot mutate them.
+    - `'runtime-immutable-def'` is a reference to a value stored in a `tgpu.const` but accessed in a way that isn’t constant at shader‑creation time (for example, indexed using a runtime value).
+    - `'local-def'` lets us track whether values originate from the function we're currently generating.
 - **Argument Origins**: This group is dedicated to exactly one origin: 'argument'. It represents values that are passed as arguments to functions.

apps/typegpu-docs/src/examples/tests/wgsl-resolution/index.ts

@@ -99,7 +99,7 @@ const mainCompute = tgpu
   })((input) => {
     const index = input.gid.x;
     const currentTrianglePos = computeBindGroupLayout.$.currentTrianglePos;
-    const instanceInfo = currentTrianglePos[index];
+    const instanceInfo = TriangleData(currentTrianglePos[index]);
     let separation = d.vec2f();
     let alignment = d.vec2f();
     let cohesion = d.vec2f();

apps/typegpu-docs/tests/individual-example-tests/wgsl-resolution.test.ts

@@ -78,7 +78,7 @@ describe('wgsl resolution example', () => {
       @compute @workgroup_size(1) fn compute_shader(@builtin(global_invocation_id) gid: vec3u) {
         let index = gid.x;
         let currentTrianglePos = (&currentTrianglePos_1);
-        let instanceInfo = (&(*currentTrianglePos)[index]);
+        var instanceInfo = (*currentTrianglePos)[index];
         var separation = vec2f();
         var alignment = vec2f();
         var cohesion = vec2f();
@@ -89,9 +89,9 @@ describe('wgsl resolution example', () => {
             continue;
           }
           let other = (&(*currentTrianglePos)[i]);
-          let dist = distance((*instanceInfo).position, (*other).position);
+          let dist = distance(instanceInfo.position, (*other).position);
           if ((dist < paramsBuffer.separationDistance)) {
-            separation = (separation + ((*instanceInfo).position - (*other).position));
+            separation = (separation + (instanceInfo.position - (*other).position));
           }
           if ((dist < paramsBuffer.alignmentDistance)) {
             alignment = (alignment + (*other).velocity);
@@ -107,27 +107,27 @@ describe('wgsl resolution example', () => {
         }
         if ((cohesionCount > 0i)) {
           cohesion = ((1f / f32(cohesionCount)) * cohesion);
-          cohesion = (cohesion - (*instanceInfo).position);
+          cohesion = (cohesion - instanceInfo.position);
         }
         var velocity = (paramsBuffer.separationStrength * separation);
         velocity = (velocity + (paramsBuffer.alignmentStrength * alignment));
         velocity = (velocity + (paramsBuffer.cohesionStrength * cohesion));
-        (*instanceInfo).velocity = ((*instanceInfo).velocity + velocity);
-        (*instanceInfo).velocity = (clamp(length((*instanceInfo).velocity), 0f, 0.01f) * normalize((*instanceInfo).velocity));
-        if (((*instanceInfo).position.x > 1.03f)) {
-          (*instanceInfo).position.x = -1.03f;
+        instanceInfo.velocity = (instanceInfo.velocity + velocity);
+        instanceInfo.velocity = (clamp(length(instanceInfo.velocity), 0f, 0.01f) * normalize(instanceInfo.velocity));
+        if ((instanceInfo.position.x > 1.03f)) {
+          instanceInfo.position.x = -1.03f;
         }
-        if (((*instanceInfo).position.y > 1.03f)) {
-          (*instanceInfo).position.y = -1.03f;
+        if ((instanceInfo.position.y > 1.03f)) {
+          instanceInfo.position.y = -1.03f;
         }
-        if (((*instanceInfo).position.x < -1.03f)) {
-          (*instanceInfo).position.x = 1.03f;
+        if ((instanceInfo.position.x < -1.03f)) {
+          instanceInfo.position.x = 1.03f;
         }
-        if (((*instanceInfo).position.y < -1.03f)) {
-          (*instanceInfo).position.y = 1.03f;
+        if ((instanceInfo.position.y < -1.03f)) {
+          instanceInfo.position.y = 1.03f;
         }
-        (*instanceInfo).position = ((*instanceInfo).position + (*instanceInfo).velocity);
-        nextTrianglePos[index] = (*instanceInfo);
+        instanceInfo.position = (instanceInfo.position + instanceInfo.velocity);
+        nextTrianglePos[index] = instanceInfo;
       }"
     `);
   });

packages/typegpu/src/core/buffer/bufferUsage.ts

@@ -1,6 +1,6 @@
 import { schemaCallWrapper } from '../../data/schemaCallWrapper.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
-import { type AnyWgslData, type BaseData, isNaturallyEphemeral } from '../../data/wgslTypes.ts';
+import { type AnyWgslData, type BaseData } from '../../data/wgslTypes.ts';
 import { IllegalBufferAccessError } from '../../errors.ts';
 import { getExecMode, inCodegenMode, isInsideTgpuFn } from '../../execMode.ts';
 import { isUsableAsStorage, type StorageFlag } from '../../extension.ts';
@@ -129,7 +129,7 @@ class TgpuFixedBufferImpl<TData extends BaseData, TUsage extends BindableBufferU
       `@group(${group}) @binding(${binding}) var<${usage}> ${id}: ${ctx.resolve(dataType).value};`,
     );
 
-    return snip(id, dataType, isNaturallyEphemeral(dataType) ? 'runtime' : this.usage);
+    return snip(id, dataType, this.usage);
   }
 
   toString(): string {
@@ -144,7 +144,7 @@ class TgpuFixedBufferImpl<TData extends BaseData, TUsage extends BindableBufferU
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(this, dataType, isNaturallyEphemeral(dataType) ? 'runtime' : usage);
+          return snip(this, dataType, usage);
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `${this.usage}:${getName(this) ?? '<unnamed>'}.$`,
@@ -251,7 +251,7 @@ export class TgpuLaidOutBufferImpl<TData extends BaseData, TUsage extends Bindab
       };`,
     );
 
-    return snip(id, this.dataType, isNaturallyEphemeral(this.dataType) ? 'runtime' : this.usage);
+    return snip(id, this.dataType, this.usage);
   }
 
   toString(): string {
@@ -266,7 +266,7 @@ export class TgpuLaidOutBufferImpl<TData extends BaseData, TUsage extends Bindab
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(this, schema, isNaturallyEphemeral(schema) ? 'runtime' : usage);
+          return snip(this, schema, usage);
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `${this.usage}:${getName(this) ?? '<unnamed>'}.$`,

packages/typegpu/src/core/constant/tgpuConstant.ts

@@ -1,11 +1,6 @@
 import { isData, type AnyData } from '../../data/dataTypes.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
-import {
-  type AnyWgslData,
-  type BaseData,
-  isNaturallyEphemeral,
-  type WgslArray,
-} from '../../data/wgslTypes.ts';
+import { type AnyWgslData, type BaseData, type WgslArray } from '../../data/wgslTypes.ts';
 import { inCodegenMode } from '../../execMode.ts';
 import type { TgpuNamable } from '../../shared/meta.ts';
 import { getName, setName } from '../../shared/meta.ts';
@@ -119,11 +114,7 @@ class TgpuConstImpl<TDataType extends BaseData> implements TgpuConst<TDataType>,
 
     ctx.addDeclaration(`const ${id}: ${resolvedDataType} = ${resolvedValue};`);
 
-    return snip(
-      id,
-      this.dataType,
-      isNaturallyEphemeral(this.dataType) ? 'constant' : 'constant-tgpu-const-ref',
-    );
+    return snip(id, this.dataType, 'constant-immutable-def');
   }
 
   toString() {
@@ -137,11 +128,7 @@ class TgpuConstImpl<TDataType extends BaseData> implements TgpuConst<TDataType>,
       {
         [$internal]: true,
         get [$ownSnippet]() {
-          return snip(
-            this,
-            dataType,
-            isNaturallyEphemeral(dataType) ? 'constant' : 'constant-tgpu-const-ref',
-          );
+          return snip(this, dataType, 'constant-immutable-def');
         },
         [$resolve]: (ctx) => ctx.resolve(this),
         toString: () => `const:${getName(this) ?? '<unnamed>'}.$`,

packages/typegpu/src/core/rawCodeSnippet/tgpuRawCodeSnippet.ts

@@ -25,13 +25,13 @@ export interface TgpuRawCodeSnippet<TDataType extends BaseData> {
 }
 
 // The origin 'function' refers to values passed in from the calling scope, which means
-// we would have access to this value anyway. Same goes for 'argument' and 'this-function',
+// we would have access to this value anyway. Same goes for 'argument' and 'local-def',
 // the values literally exist in the function we're writing.
 //
-// 'constant-ref' was excluded because it's a special origin reserved for tgpu.const values.
+// '*-immutable-def' were excluded because they're a special origin reserved for tgpu.const values.
 export type RawCodeSnippetOrigin = Exclude<
   Origin,
-  'function' | 'this-function' | 'argument' | 'constant-ref'
+  'function' | 'local-def' | 'argument' | 'constant-immutable-def' | 'runtime-immutable-def'
 >;
 
 /**

packages/typegpu/src/core/variable/tgpuVariable.ts

@@ -1,6 +1,6 @@
 import type { AnyData } from '../../data/dataTypes.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
-import { type BaseData, isNaturallyEphemeral } from '../../data/wgslTypes.ts';
+import type { BaseData } from '../../data/wgslTypes.ts';
 import { IllegalVarAccessError } from '../../errors.ts';
 import { getExecMode, isInsideTgpuFn } from '../../execMode.ts';
 import type { TgpuNamable } from '../../shared/meta.ts';
@@ -96,7 +96,7 @@ class TgpuVarImpl<TScope extends VariableScope, TDataType extends BaseData>
       ctx.addDeclaration(`${pre};`);
     }
 
-    return snip(id, this.#dataType, isNaturallyEphemeral(this.#dataType) ? 'runtime' : this.#scope);
+    return snip(id, this.#dataType, this.#scope);
   }
 
   $name(label: string) {
@@ -110,7 +110,7 @@ class TgpuVarImpl<TScope extends VariableScope, TDataType extends BaseData>
 
   get [$gpuValueOf](): InferGPU<TDataType> {
     const dataType = this.#dataType;
-    const origin = isNaturallyEphemeral(dataType) ? 'runtime' : this.#scope;
+    const origin = this.#scope;
 
     return new Proxy(
       {

packages/typegpu/src/data/ref.ts

@@ -5,7 +5,7 @@ import { $gpuCallable, $internal, $ownSnippet, $resolve } from '../shared/symbol
 import type { DualFn, SelfResolvable } from '../types.ts';
 import { UnknownData } from './dataTypes.ts';
 import { createPtrFromOrigin, explicitFrom } from './ptr.ts';
-import { type ResolvedSnippet, snip, type Snippet } from './snippet.ts';
+import { isAlias, type ResolvedSnippet, snip, type Snippet } from './snippet.ts';
 import { isNaturallyEphemeral, isPtr, type Ptr, type StorableData } from './wgslTypes.ts';
 
 // ----------
@@ -47,10 +47,30 @@ export const _ref = (() => {
   impl.toString = () => 'ref';
   impl[$internal] = true;
   impl[$gpuCallable] = {
-    call(_ctx, [value]) {
+    call(ctx, [value]) {
       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.`,
+          stitch`d.ref(${value}) is illegal, cannot take a reference of an argument. Copy the value first, and take a reference of the copy.`,
+        );
+      }
+
+      if (value.origin === 'constant-immutable-def' || value.origin === 'runtime-immutable-def') {
+        const typeStr = ctx.resolve(value.dataType).value;
+        throw new WgslTypeError(
+          stitch`d.ref(${value}) is illegal, cannot take a reference to a constant.
+-----
+- Try 'd.ref(${typeStr}(${value}));' instead to create a new referencable value.
+-----`,
+        );
+      }
+
+      if (isAlias(value) && isNaturallyEphemeral(value.dataType)) {
+        const typeStr = ctx.resolve(value.dataType).value;
+        throw new WgslTypeError(
+          stitch`d.ref(${value}) is illegal, cannot take a reference to a scalar value.
+-----
+- Try 'd.ref(${typeStr}(${value}));' instead to create a new referencable scalar.
+-----`,
         );
       }
 
@@ -172,11 +192,10 @@ export function derefSnippet(snippet: Snippet): Snippet {
   }
 
   const innerType = snippet.dataType.inner;
-  const origin = isNaturallyEphemeral(innerType) ? 'runtime' : snippet.origin;
 
   if (snippet.value instanceof RefOperator) {
-    return snip(stitch`${snippet.value.snippet}`, innerType, origin);
+    return snip(stitch`${snippet.value.snippet}`, innerType, snippet.origin);
   }
 
-  return snip(stitch`(*${snippet})`, innerType, origin);
+  return snip(stitch`(*${snippet})`, innerType, snippet.origin);
 }

packages/typegpu/src/data/snippet.ts

@@ -4,34 +4,34 @@ import { DEV } from '../shared/env.ts';
 import { type BaseData, isNumericSchema } from './wgslTypes.ts';
 
 export type Origin =
-  | 'uniform'
-  | 'readonly' // equivalent to ptr<storage, ..., read>
-  | 'mutable' // equivalent to ptr<storage, ..., read-write>
-  | 'workgroup'
-  | 'private'
-  | 'function'
-  | 'this-function'
-  | 'handle'
-  // is an argument (or part of an argument) given to the
-  // function we're resolving. This includes primitives, to
-  // catch cases where we update an argument's primitive member
-  // prop, e.g.: `vec.x += 1;`
+  // --- ADDRESS SPACE ORIGINS
+  | 'uniform' /*   defined in the 'uniform' address space  */
+  | 'readonly' /*  defined in the 'storage' address space, with 'read' access  */
+  | 'mutable' /*   defined in the 'storage' address space, with 'read-write' access  */
+  | 'workgroup' /* defined in the 'workgroup' address space  */
+  | 'private' /*   defined in the 'private' address space  */
+  | 'handle' /*    defined in the 'handle' address space  */
+  | 'function' /*  defined in a callee, passed down to us as an argument ('function' address space)  */
+  // --- DEFINITIONS
+  // defined in the current function
+  | 'local-def'
+  // A reference to a deeply immutable definition, recognized by WGSL as a 'constant'.
+  // This is the usual case, read about 'runtime-immutable-def' to know when this doesn't apply.
+  // A reference to a tgpu.const().$ value (which is frozen) fits into this category.
+  | 'constant-immutable-def'
+  // A reference to a deeply immutable definition, NOT recognized by WGSL as a 'constant'.
+  // This can happen if say a constant is accessed with a runtime-known index. WGSL doesn't treat it like
+  // a 'constant' anymore, but it's still a frozen value in JS, so we must treat it like it's immutable.
+  | 'runtime-immutable-def'
+  // ---------
+  // non-pointer function arguments (or part of an argument).
   | 'argument'
-  // not a ref to anything, known at runtime
+  // not a reference to anything, known at runtime
   | 'runtime'
-  // not a ref to anything, known at pipeline creation time
-  // (not to be confused with 'comptime')
-  // note that this doesn't automatically mean the value can be stored in a `const`
-  // variable, more so that it's valid to do so in WGSL (but not necessarily safe to do in JS shaders)
-  | 'constant'
-  // don't even get me started on these. They're references to non-primitive values that originate
-  // from a tgpu.const(...).$ call.
-  | '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';
-}
+  // an ephemeral value that is a valid WGSL 'constant' (not to be confused with 'comptime')
+  // doesn't always lead to creating a `const` variable, as we cannot always guarantee that
+  // the value won't be mutated in JS
+  | 'constant';
 
 /**
  * What happens to a snippet's origin when it's deep copied in JS, and left as is in WGSL?
@@ -49,8 +49,27 @@ export function fallthroughCopyOrigin(origin: Origin): Origin {
   return 'runtime';
 }
 
-export function isEphemeralSnippet(snippet: Snippet) {
-  return isEphemeralOrigin(snippet.origin);
+/**
+ * Whether a snippet aliases a value that lives outside the current expression.
+ *
+ * @example
+ * ```ts
+ * function foo(a: number) {
+ *   const color = d.vec3f(1, 2, 3);
+ *   return color * a;
+ * }
+ *
+ * // References:
+ * // -  color
+ * // -  a
+ * //
+ * // Not references:
+ * // - d.vec3f(1, 2, 3)
+ * // - color * a
+ * ```
+ */
+export function isAlias(snippet: Snippet) {
+  return !(snippet.origin === 'runtime' || snippet.origin === 'constant');
 }
 
 export const originToPtrParams = {
@@ -60,7 +79,8 @@ export const originToPtrParams = {
   workgroup: { space: 'workgroup', access: 'read-write' },
   private: { space: 'private', access: 'read-write' },
   function: { space: 'function', access: 'read-write' },
-  'this-function': { space: 'function', access: 'read-write' },
+  // Local declarations are also in the `function` address space
+  'local-def': { space: 'function', access: 'read-write' },
 } as const;
 export type OriginToPtrParams = typeof originToPtrParams;