fix: ArrayExpression - immediate index access

Author: iwoplaza

packages/typegpu/src/tgsl/accessIndex.ts

@@ -13,7 +13,7 @@ import {
 } from '../data/wgslTypes.ts';
 import { isKnownAtComptime } from '../types.ts';
 import { accessProp } from './accessProp.ts';
-import { coerceToSnippet } from './generationHelpers.ts';
+import { ArrayExpression, coerceToSnippet } from './generationHelpers.ts';
 
 const indexableTypeToResult = {
   mat2x2f: vec2f,
@@ -54,6 +54,10 @@ export function accessIndex(target: Snippet, indexArg: Snippet | number): Snippe
       origin = 'runtime';
     }
 
+    if (target.value instanceof ArrayExpression && isKnownAtComptime(index)) {
+      return target.value.elements[index.value as number];
+    }
+
     return snip(
       isKnownAtComptime(target) && isKnownAtComptime(index)
         ? // oxlint-disable-next-line typescript/no-explicit-any -- it's fine, it's there

packages/typegpu/tests/array.test.ts

@@ -443,6 +443,38 @@ describe('array', () => {
       }"
     `);
   });
+
+  it('array expressions can be indexed into with a comptime-known index', () => {
+    function foo() {
+      'use gpu';
+      const i = 2;
+      const a = [i, 2][0];
+      const b = [i, 2][1];
+    }
+
+    expect(tgpu.resolve([foo])).toMatchInlineSnapshot(`
+      "fn foo() {
+        const i = 2;
+        const a = i;
+        const b = 2i;
+      }"
+    `);
+  });
+
+  it('array expressions can be indexed into with a runtime-known index', () => {
+    function foo() {
+      'use gpu';
+      const i = 0;
+      const a = [1, 2][i];
+    }
+
+    expect(tgpu.resolve([foo])).toMatchInlineSnapshot(`
+      "fn foo() {
+        const i = 0;
+        let a = array<i32, 2>(1, 2)[i];
+      }"
+    `);
+  });
 });
 
 describe('array.length', () => {