fix: Indexing constants with runtime indices should properly adjust

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

@@ -200,6 +200,7 @@ There are essentially three types of origins:
   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'`.
     - `'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-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.
+    - `'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 an expression that still references a `tgpu.const` value, but isn't considered constant in the eyes of WGSL.
     - `'this-function'` lets us track whether values originates from the function we're currently generating, or the function that called us.
 - **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/constant/tgpuConstant.ts

@@ -98,7 +98,9 @@ class TgpuConstImpl<TDataType extends AnyWgslData>
     return snip(
       id,
       this.dataType,
-      isNaturallyEphemeral(this.dataType) ? 'constant' : 'constant-ref',
+      isNaturallyEphemeral(this.dataType)
+        ? 'constant'
+        : 'constant-tgpu-const-ref',
     );
   }
 
@@ -115,7 +117,9 @@ class TgpuConstImpl<TDataType extends AnyWgslData>
         return snip(
           this,
           dataType,
-          isNaturallyEphemeral(dataType) ? 'constant' : 'constant-ref',
+          isNaturallyEphemeral(dataType)
+            ? 'constant'
+            : 'constant-tgpu-const-ref',
         );
       },
       [$resolve]: (ctx) => ctx.resolve(this),

packages/typegpu/src/data/snippet.ts

@@ -24,7 +24,10 @@ export type Origin =
   // 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'
-  | 'constant-ref';
+  // 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';

packages/typegpu/src/tgsl/generationHelpers.ts

@@ -20,6 +20,7 @@ import {
 import {
   isEphemeralSnippet,
   isSnippet,
+  type Origin,
   snip,
   type Snippet,
 } from '../data/snippet.ts';
@@ -156,10 +157,7 @@ export function accessProp(
   target: Snippet,
   propName: string,
 ): Snippet | undefined {
-  if (
-    infixKinds.includes(target.dataType.type) &&
-    propName in infixOperators
-  ) {
+  if (infixKinds.includes(target.dataType.type) && propName in infixOperators) {
     return snip(
       new InfixDispatch(
         propName,
@@ -174,11 +172,7 @@ export function accessProp(
   if (isWgslArray(target.dataType) && propName === 'length') {
     if (target.dataType.elementCount === 0) {
       // Dynamically-sized array
-      return snip(
-        stitch`arrayLength(&${target})`,
-        u32,
-        /* origin */ 'runtime',
-      );
+      return snip(stitch`arrayLength(&${target})`, u32, /* origin */ 'runtime');
     }
 
     return snip(
@@ -208,10 +202,10 @@ export function accessProp(
       propType,
       /* origin */ target.origin === 'argument'
         ? 'argument'
-        : !isEphemeralSnippet(target) &&
-            !isNaturallyEphemeral(propType)
+        : !isEphemeralSnippet(target) && !isNaturallyEphemeral(propType)
         ? target.origin
-        : target.origin === 'constant' || target.origin === 'constant-ref'
+        : target.origin === 'constant' ||
+            target.origin === 'constant-tgpu-const-ref'
         ? 'constant'
         : 'runtime',
     );
@@ -231,11 +225,7 @@ export function accessProp(
   }
 
   const propLength = propName.length;
-  if (
-    isVec(target.dataType) &&
-    propLength >= 1 &&
-    propLength <= 4
-  ) {
+  if (isVec(target.dataType) && propLength >= 1 && propLength <= 4) {
     const swizzleTypeChar = target.dataType.type.includes('bool')
       ? 'b'
       : (target.dataType.type[4] as SwizzleableType);
@@ -255,7 +245,7 @@ export function accessProp(
       /* origin */ target.origin === 'argument' && propLength === 1
         ? 'argument'
         : target.origin === 'constant' ||
-            target.origin === 'constant-ref'
+            target.origin === 'constant-tgpu-const-ref'
         ? 'constant'
         : 'runtime',
     );
@@ -282,19 +272,36 @@ export function accessIndex(
   // array
   if (isWgslArray(target.dataType) || isDisarray(target.dataType)) {
     const elementType = target.dataType.elementType as AnyData;
+    const isTargetEphemeral = isEphemeralSnippet(target);
+    const isElementNatEph = isNaturallyEphemeral(elementType);
+
+    let origin: Origin;
+
+    if (
+      target.origin === 'constant-tgpu-const-ref' &&
+      index.origin === 'constant'
+    ) {
+      origin = isElementNatEph ? 'constant' : 'constant-tgpu-const-ref';
+    } else if (
+      target.origin === 'constant-tgpu-const-ref' ||
+      target.origin === 'runtime-tgpu-const-ref'
+    ) {
+      origin = isElementNatEph ? 'runtime' : 'runtime-tgpu-const-ref';
+    } else if (!isTargetEphemeral && !isElementNatEph) {
+      origin = target.origin;
+    } else if (target.origin === 'constant' && index.origin === 'constant') {
+      origin = 'constant';
+    } else {
+      origin = 'runtime';
+    }
 
     return snip(
       isKnownAtComptime(target) && isKnownAtComptime(index)
         // biome-ignore lint/suspicious/noExplicitAny: it's fine, it's there
         ? (target.value as any)[index.value as number]
         : stitch`${target}[${index}]`,
       elementType,
-      /* origin */ !isEphemeralSnippet(target) &&
-          !isNaturallyEphemeral(elementType)
-        ? target.origin
-        : target.origin === 'constant' || target.origin === 'constant-ref'
-        ? 'constant'
-        : 'runtime',
+      /* origin */ origin,
     );
   }
 
@@ -307,7 +314,7 @@ export function accessIndex(
         : stitch`${target}[${index}]`,
       target.dataType.primitive,
       /* origin */ target.origin === 'constant' ||
-          target.origin === 'constant-ref'
+          target.origin === 'constant-tgpu-const-ref'
         ? 'constant'
         : 'runtime',
     );

packages/typegpu/src/tgsl/shellless.ts

@@ -67,7 +67,10 @@ export class ShelllessRepository {
 
       let type = concretize(s.dataType as AnyData);
 
-      if (s.origin === 'constant-ref') {
+      if (
+        s.origin === 'constant-tgpu-const-ref' ||
+        s.origin === 'runtime-tgpu-const-ref'
+      ) {
         // biome-ignore lint/style/noNonNullAssertion: it's there
         const ctx = getResolutionCtx()!;
         throw new Error(

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -172,10 +172,13 @@ ${this.ctx.pre}}`;
     const naturallyEphemeral = wgsl.isNaturallyEphemeral(dataType);
 
     let varOrigin: Origin;
-    if (origin === 'constant-ref') {
+    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 = 'constant-ref';
+      varOrigin = origin;
     } else if (origin === 'argument') {
       if (naturallyEphemeral) {
         varOrigin = 'runtime';
@@ -288,7 +291,9 @@ ${this.ctx.pre}}`;
 
       if (exprType === NODE.assignmentExpr) {
         if (
-          convLhs.origin === 'constant' || convLhs.origin === 'constant-ref'
+          convLhs.origin === 'constant' ||
+          convLhs.origin === 'constant-tgpu-const-ref' ||
+          convLhs.origin === 'runtime-tgpu-const-ref'
         ) {
           throw new WgslTypeError(
             `'${lhsStr} = ${rhsStr}' is invalid, because ${lhsStr} is a constant.`,
@@ -875,8 +880,10 @@ ${this.ctx.pre}else ${alternate}`;
           );
         }
 
-        if (eq.origin === 'constant-ref') {
+        if (eq.origin === 'constant-tgpu-const-ref') {
           varType = 'const';
+        } else if (eq.origin === 'runtime-tgpu-const-ref') {
+          varType = 'let';
         } else {
           varType = 'let';
           if (!wgsl.isPtr(dataType)) {

packages/typegpu/tests/constant.test.ts

@@ -91,4 +91,34 @@ describe('tgpu.const', () => {
       `[TypeError: Cannot assign to read only property 'pos' of object '#<Object>']`,
     );
   });
+
+  it('looses its `constant` origin when indexing with runtime value', () => {
+    const positions = tgpu.const(d.arrayOf(d.vec3f, 3), [
+      d.vec3f(0),
+      d.vec3f(1),
+      d.vec3f(2),
+    ]);
+
+    const foo = (idx: number) => {
+      'use gpu';
+      const pos = positions.$[idx];
+    };
+
+    const main = () => {
+      'use gpu';
+      foo(0);
+    };
+
+    expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
+      "const positions: array<vec3f, 3> = array<vec3f, 3>(vec3f(), vec3f(1), vec3f(2));
+
+      fn foo(idx: i32) {
+        let pos = positions[idx];
+      }
+
+      fn main() {
+        foo(0i);
+      }"
+    `);
+  });
 });