refactor: Better origin tracking

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.
+    - `'local'` 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.

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-tgpu-const-ref');
   }
 
   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-tgpu-const-ref');
         },
         [$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',
 // 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.
+// '*-tgpu-const-ref' 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' | 'argument' | 'constant-tgpu-const-ref' | 'runtime-tgpu-const-ref'
 >;
 
 /**

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

@@ -6,7 +6,7 @@ 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 { isNaturallyEphemeral, isPtr, type Ptr, type StorableData } from './wgslTypes.ts';
+import { isPtr, type Ptr, type StorableData } from './wgslTypes.ts';
 
 // ----------
 // Public API
@@ -172,11 +172,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

@@ -1,37 +1,39 @@
 import { undecorate } from './dataTypes.ts';
 import type { UnknownData } from './dataTypes.ts';
 import { DEV } from '../shared/env.ts';
-import { type BaseData, isNumericSchema } from './wgslTypes.ts';
+import { type BaseData, isNaturallyEphemeral, 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)  */
+  // ---------
+  // non-pointer function arguments (or part of an argument).
   | 'argument'
-  // not a ref to anything, known at runtime
+  // defined in the current function
+  | 'local'
+  // 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)
+  // 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'
-  // 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';
-}
+  //
+  // TGPU.CONST REFERENCES
+  // A reference to a tgpu.const().$ value (which is frozen), that is treated by WGSL as a 'constant'.
+  // This is the usual case, read about 'runtime-tgpu-const-ref' to know when this doesn't apply.
+  // Turns into a `const` when assigned to a variable
+  | 'constant-tgpu-const-ref'
+  // A reference to a tgpu.const().$ value (which is frozen), that is NOT treated 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.
+  // Turns into a `let` when assigned to a variable
+  | 'runtime-tgpu-const-ref';
 
 /**
  * What happens to a snippet's origin when it's deep copied in JS, and left as is in WGSL?
@@ -49,8 +51,29 @@ export function fallthroughCopyOrigin(origin: Origin): Origin {
   return 'runtime';
 }
 
-export function isEphemeralSnippet(snippet: Snippet) {
-  return isEphemeralOrigin(snippet.origin);
+export function isReference(snippet: Snippet) {
+  return !(snippet.origin === 'runtime' || snippet.origin === 'constant');
+}
+
+/**
+ * TODO: Remove this in favor of just writing plain logic. This removes the need to name this behavior in any way
+ *
+ * Snippets that "carry mutable references" are those that, when assigned to a variable,
+ * assign a name to a value that is already mutable by other means in the shader.
+ *
+ * ```ts
+ * // "d.vec3f()" is a fresh value, doesn't carry a reference to anything already existing before
+ * const vec = d.vec3f();
+ * const foo = vec; // the vector now is mutable through `vec` and accessible through `foo`.
+ * ```
+ */
+export function assignmentIsReassignment(snippet: Snippet): boolean {
+  if (isNaturallyEphemeral(snippet.dataType)) {
+    return false;
+  }
+  if (snippet.origin === 'function' && snippet.dataType) {
+  }
+  return !(snippet.origin === 'runtime' || snippet.origin === 'constant');
 }
 
 export const originToPtrParams = {
@@ -60,7 +83,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: { space: 'function', access: 'read-write' },
 } as const;
 export type OriginToPtrParams = typeof originToPtrParams;
 

packages/typegpu/src/tgsl/accessIndex.ts

@@ -1,16 +1,9 @@
 import { stitch } from '../core/resolve/stitch.ts';
 import { isDisarray, MatrixColumnsAccess } from '../data/dataTypes.ts';
 import { derefSnippet } from '../data/ref.ts';
-import { isEphemeralSnippet, type Origin, snip, type Snippet } from '../data/snippet.ts';
+import { type Origin, snip, type Snippet } from '../data/snippet.ts';
 import { vec2f, vec3f, vec4f } from '../data/vector.ts';
-import {
-  type BaseData,
-  isNaturallyEphemeral,
-  isPtr,
-  isVec,
-  isWgslArray,
-  isWgslStruct,
-} from '../data/wgslTypes.ts';
+import { type BaseData, isPtr, isVec, isWgslArray, isWgslStruct } from '../data/wgslTypes.ts';
 import { isKnownAtComptime } from '../types.ts';
 import { accessProp } from './accessProp.ts';
 import { ArrayExpression, coerceToSnippet } from './generationHelpers.ts';
@@ -27,31 +20,18 @@ export function accessIndex(target: Snippet, indexArg: Snippet | number): Snippe
   // array
   if (isWgslArray(target.dataType) || isDisarray(target.dataType)) {
     const elementType = target.dataType.elementType;
-    const isElementNatEph = isNaturallyEphemeral(elementType);
-    const isTargetEphemeral = isEphemeralSnippet(target);
-    const isIndexConstant = index.origin === 'constant';
 
     let origin: Origin;
 
     if (target.origin === 'constant-tgpu-const-ref') {
       // Constant refs stay const unless the element/index forces runtime materialization
-      if (isIndexConstant) {
-        origin = isElementNatEph ? 'constant' : 'constant-tgpu-const-ref';
-      } else {
-        origin = isElementNatEph ? 'runtime' : 'runtime-tgpu-const-ref';
-      }
-    } else if (target.origin === 'runtime-tgpu-const-ref') {
-      // Runtime refs keep their ref semantics unless the element is ephemeral only
-      origin = isElementNatEph ? 'runtime' : 'runtime-tgpu-const-ref';
-    } else if (!isTargetEphemeral && !isElementNatEph) {
-      // Stable containers can forward their origin information
-      origin = target.origin;
-    } else if (isIndexConstant && target.origin === 'constant') {
-      // Plain constants indexed with constants stay constant
-      origin = 'constant';
+      origin = index.origin === 'constant' ? 'constant-tgpu-const-ref' : 'runtime-tgpu-const-ref';
+    } else if (target.origin === 'constant') {
+      // Ephemeral constants indexed with constants stay constant, otherwise they become runtime-known
+      origin = index.origin === 'constant' ? 'constant' : 'runtime';
     } else {
-      // Everything else must be produced at runtime
-      origin = 'runtime';
+      // Fallthrough
+      origin = target.origin;
     }
 
     if (target.value instanceof ArrayExpression && isKnownAtComptime(index)) {
@@ -76,9 +56,7 @@ export function accessIndex(target: Snippet, indexArg: Snippet | number): Snippe
           (target.value as any)[index.value as any]
         : stitch`${target}[${index}]`,
       target.dataType.primitive,
-      /* origin */ target.origin === 'constant' || target.origin === 'constant-tgpu-const-ref'
-        ? 'constant'
-        : 'runtime',
+      /* origin */ target.origin,
     );
   }
 

packages/typegpu/src/tgsl/accessProp.ts

@@ -10,7 +10,7 @@ import {
 } from '../data/dataTypes.ts';
 import { abstractInt, bool, f16, f32, i32, u32 } from '../data/numeric.ts';
 import { derefSnippet } from '../data/ref.ts';
-import { isEphemeralSnippet, isSnippet, snip, type Snippet } from '../data/snippet.ts';
+import { isSnippet, snip, type Snippet } from '../data/snippet.ts';
 import {
   vec2b,
   vec2f,
@@ -31,7 +31,6 @@ import {
 import {
   type BaseData,
   isMat,
-  isNaturallyEphemeral,
   isPtr,
   isVec,
   isWgslArray,
@@ -138,17 +137,7 @@ export function accessProp(target: Snippet, propName: string): Snippet | undefin
     }
     propType = undecorate(propType);
 
-    return snip(
-      stitch`${target}.${propName}`,
-      propType,
-      /* origin */ target.origin === 'argument'
-        ? 'argument'
-        : !isEphemeralSnippet(target) && !isNaturallyEphemeral(propType)
-          ? target.origin
-          : target.origin === 'constant' || target.origin === 'constant-tgpu-const-ref'
-            ? 'constant'
-            : 'runtime',
-    );
+    return snip(stitch`${target}.${propName}`, propType, /* origin */ target.origin);
   }
 
   if (target.dataType instanceof AutoStruct) {
@@ -215,9 +204,9 @@ export function accessProp(target: Snippet, propName: string): Snippet | undefin
         : stitch`${target}.${propName}`,
       swizzleType,
       // Swizzling creates new vectors (unless they're on the lhs of an assignment, but that's not yet supported in WGSL)
-      /* origin */ target.origin === 'argument' && propLength === 1
-        ? 'argument'
-        : target.origin === 'constant' || target.origin === 'constant-tgpu-const-ref'
+      /* origin */ propLength === 1
+        ? target.origin
+        : target.origin === 'constant'
           ? 'constant'
           : 'runtime',
     );