fix: Throw error when accessing matrix elements directly (#1625)

Author: aleksanderkatan

packages/typegpu/src/data/dataTypes.ts

@@ -236,3 +236,9 @@ export class InfixDispatch {
     readonly operator: (lhs: Snippet, rhs: Snippet) => Snippet,
   ) {}
 }
+
+export class MatrixColumnsAccess {
+  constructor(
+    readonly matrix: Snippet,
+  ) {}
+}

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -1,19 +1,27 @@
 import * as tinyest from 'tinyest';
+import { stitch, stitchWithExactTypes } from '../core/resolve/stitch.ts';
 import { arrayOf } from '../data/array.ts';
 import {
   type AnyData,
   InfixDispatch,
   isData,
   isLooseData,
+  MatrixColumnsAccess,
   UnknownData,
 } from '../data/dataTypes.ts';
-import { isSnippet, snip, type Snippet } from '../data/snippet.ts';
 import { abstractInt, bool, u32 } from '../data/numeric.ts';
+import { isSnippet, snip, type Snippet } from '../data/snippet.ts';
 import * as wgsl from '../data/wgslTypes.ts';
 import { ResolutionError, WgslTypeError } from '../errors.ts';
 import { getName } from '../shared/meta.ts';
 import { $internal } from '../shared/symbols.ts';
+import { add, div, mul, sub } from '../std/operators.ts';
 import { type FnArgsConversionHint, isMarkedInternal } from '../types.ts';
+import {
+  convertStructValues,
+  convertToCommonType,
+  tryConvertSnippet,
+} from './conversion.ts';
 import {
   coerceToSnippet,
   concretize,
@@ -22,13 +30,6 @@ import {
   getTypeForPropAccess,
   numericLiteralToSnippet,
 } from './generationHelpers.ts';
-import {
-  convertStructValues,
-  convertToCommonType,
-  tryConvertSnippet,
-} from './conversion.ts';
-import { add, div, mul, sub } from '../std/operators.ts';
-import { stitch, stitchWithExactTypes } from '../core/resolve/stitch.ts';
 
 const { NodeTypeCatalog: NODE } = tinyest;
 
@@ -309,7 +310,7 @@ export function generateExpression(
     }
 
     if (wgsl.isMat(target.dataType) && property === 'columns') {
-      return snip(target.value, target.dataType);
+      return snip(new MatrixColumnsAccess(target), UnknownData);
     }
 
     if (
@@ -331,9 +332,16 @@ export function generateExpression(
     const [_, targetNode, propertyNode] = expression;
     const target = generateExpression(ctx, targetNode);
     const property = generateExpression(ctx, propertyNode);
-    const targetStr = ctx.resolve(target.value, target.dataType);
     const propertyStr = ctx.resolve(property.value, property.dataType);
 
+    if (target.value instanceof MatrixColumnsAccess) {
+      return snip(
+        stitch`${target.value.matrix}[${propertyStr}]`,
+        getTypeForIndexAccess(target.value.matrix.dataType as AnyData),
+      );
+    }
+    const targetStr = ctx.resolve(target.value, target.dataType);
+
     if (target.dataType.type === 'unknown') {
       // No idea what the type is, so we act on the snippet's value and try to guess
 
@@ -351,6 +359,12 @@ export function generateExpression(
       );
     }
 
+    if (wgsl.isMat(target.dataType)) {
+      throw new Error(
+        "The only way of accessing matrix elements in TGSL is through the 'columns' property.",
+      );
+    }
+
     if (wgsl.isPtr(target.dataType)) {
       return snip(
         `(*${targetStr})[${propertyStr}]`,

packages/typegpu/tests/examples/individual/fluid-with-atomics.test.ts

@@ -21,24 +21,24 @@ describe('fluid with atomics example', () => {
         @builtin(global_invocation_id) gid: vec3u,
       }
 
-      @group(0) @binding(0) var<storage, read_write> nextState_5: array<atomic<u32>, 1048576>;
+      @group(0) @binding(0) var<uniform> size_6: vec2u;
 
-      @group(0) @binding(1) var<uniform> size_7: vec2u;
-
-      fn getIndex_6(x: u32, y: u32) -> u32 {
-        var h = size_7.y;
-        var w = size_7.x;
+      fn getIndex_5(x: u32, y: u32) -> u32 {
+        var h = size_6.y;
+        var w = size_6.x;
         return (((y % h) * w) + (x % w));
       }
 
+      @group(0) @binding(1) var<storage, read_write> nextState_7: array<atomic<u32>, 1048576>;
+
       fn updateCell_4(x: u32, y: u32, value: u32) {
-        atomicStore(&nextState_5[getIndex_6(x, y)], value);
+        atomicStore(&nextState_7[getIndex_5(x, y)], value);
       }
 
       @group(0) @binding(2) var<storage, read> currentStateBuffer_10: array<u32, 1048576>;
 
       fn getCell_9(x: u32, y: u32) -> u32 {
-        return currentStateBuffer_10[getIndex_6(x, y)];
+        return currentStateBuffer_10[getIndex_5(x, y)];
       }
 
       fn isClearCell_8(x: u32, y: u32) -> bool {
@@ -59,14 +59,14 @@ describe('fluid with atomics example', () => {
       }
 
       fn getCellNext_15(x: u32, y: u32) -> u32 {
-        return atomicLoad(&nextState_5[getIndex_6(x, y)]);
+        return atomicLoad(&nextState_7[getIndex_5(x, y)]);
       }
 
       fn addToCell_14(x: u32, y: u32, value: u32) {
         var cell = getCellNext_15(x, y);
         var waterLevel = (cell & MAX_WATER_LEVEL_12);
         var newWaterLevel = min((waterLevel + value), MAX_WATER_LEVEL_12);
-        atomicAdd(&nextState_5[getIndex_6(x, y)], (newWaterLevel - waterLevel));
+        atomicAdd(&nextState_7[getIndex_5(x, y)], (newWaterLevel - waterLevel));
       }
 
       fn isWaterSource_16(x: u32, y: u32) -> bool {
@@ -81,7 +81,7 @@ describe('fluid with atomics example', () => {
         var cell = getCellNext_15(x, y);
         var waterLevel = (cell & MAX_WATER_LEVEL_12);
         var newWaterLevel = max((waterLevel - min(value, waterLevel)), 0);
-        atomicSub(&nextState_5[getIndex_6(x, y)], (waterLevel - newWaterLevel));
+        atomicSub(&nextState_7[getIndex_5(x, y)], (waterLevel - newWaterLevel));
       }
 
       fn getWaterLevel_19(x: u32, y: u32) -> u32 {
@@ -107,7 +107,7 @@ describe('fluid with atomics example', () => {
           updateCell_4(x, y, (3 << 24));
           return true;
         }
-        if (((((y == 0) || (y == (size_7.y - 1))) || (x == 0)) || (x == (size_7.x - 1)))) {
+        if (((((y == 0) || (y == (size_6.y - 1))) || (x == 0)) || (x == (size_6.x - 1)))) {
           subtractFromCell_18(x, y, getWaterLevel_19(x, y));
           return true;
         }

packages/typegpu/tests/tgsl/wgslGenerator.test.ts

@@ -1,8 +1,8 @@
 import * as tinyest from 'tinyest';
 import { beforeEach, describe, expect } from 'vitest';
-import { snip } from '../../src/data/snippet.ts';
 import * as d from '../../src/data/index.ts';
 import { abstractFloat, abstractInt } from '../../src/data/numeric.ts';
+import { snip } from '../../src/data/snippet.ts';
 import { Void, type WgslArray } from '../../src/data/wgslTypes.ts';
 import { provideCtx } from '../../src/execMode.ts';
 import tgpu from '../../src/index.ts';
@@ -932,4 +932,55 @@ describe('wgslGenerator', () => {
       }`),
     );
   });
+
+  it('throws error when accessing matrix elements directly', () => {
+    const testFn = tgpu.fn([])(() => {
+      const matrix = d.mat4x4f();
+      const element = matrix[4];
+    });
+
+    expect(() => asWgsl(testFn))
+      .toThrowErrorMatchingInlineSnapshot(`
+        [Error: Resolution of the following tree failed:
+        - <root>
+        - fn:testFn: The only way of accessing matrix elements in TGSL is through the 'columns' property.]
+      `);
+  });
+
+  it('generates correct code when accessing matrix elements through .columns', () => {
+    const testFn = tgpu.fn([])(() => {
+      const matrix = d.mat4x4f();
+      const column = matrix.columns[1];
+      const element = column[0];
+      const directElement = matrix.columns[1][0];
+    });
+
+    expect(asWgsl(testFn)).toMatchInlineSnapshot(`
+      "fn testFn() {
+        var matrix = mat4x4f();
+        var column = matrix[1];
+        var element = column[0];
+        var directElement = matrix[1][0];
+      }"
+    `);
+  });
+
+  it('resolves when accessing matrix elements through .columns', () => {
+    const matrix = tgpu['~unstable'].workgroupVar(d.mat4x4f);
+    const index = tgpu['~unstable'].workgroupVar(d.u32);
+
+    const testFn = tgpu.fn([])(() => {
+      const element = matrix.$.columns[index.$];
+    });
+
+    expect(asWgsl(testFn)).toMatchInlineSnapshot(`
+      "var<workgroup> index: u32;
+
+      var<workgroup> matrix: mat4x4f;
+
+      fn testFn() {
+        var element = matrix[index];
+      }"
+    `);
+  });
 });