feat: array.length in TGSL (#1097)

Author: iwoplaza

apps/typegpu-docs/src/content/examples/simulation/boids-next/index.ts

@@ -30,11 +30,10 @@ const TriangleData = d.struct({
 });
 
 const renderBindGroupLayout = tgpu.bindGroupLayout({
-  trianglePos: { uniform: d.arrayOf(TriangleData, triangleAmount) },
   colorPalette: { uniform: d.vec3f },
 });
 
-const { trianglePos, colorPalette } = renderBindGroupLayout.bound;
+const { colorPalette } = renderBindGroupLayout.bound;
 
 const VertexOutput = {
   position: d.builtin.position,
@@ -61,7 +60,6 @@ const mainVert = tgpu['~unstable']
     return VertexOutput(pos, color);
   }`)
   .$uses({
-    trianglePos,
     colorPalette,
     getRotationFromVelocity,
     rotate,
@@ -223,7 +221,7 @@ const mainCompute = tgpu['~unstable']
     let alignmentCount = 0;
     let cohesionCount = 0;
 
-    for (let i = d.u32(0); i < std.arrayLength(currentTrianglePos.value); i++) {
+    for (let i = d.u32(0); i < currentTrianglePos.value.length; i++) {
       if (i === index) {
         continue;
       }
@@ -295,7 +293,6 @@ const computePipeline = root['~unstable']
 
 const renderBindGroups = [0, 1].map((idx) =>
   root.createBindGroup(renderBindGroupLayout, {
-    trianglePos: trianglePosBuffers[idx],
     colorPalette: colorPaletteBuffer,
   }),
 );

packages/typegpu/src/smol/wgslGenerator.ts

@@ -1,19 +1,20 @@
 import type * as smol from 'tinyest';
 import * as d from '../data';
-import * as wgsl from '../data/wgslTypes';
+import { abstractInt } from '../data/numeric.js';
+import * as wgsl from '../data/wgslTypes.js';
 import {
   type ResolutionCtx,
   type Resource,
   UnknownData,
   type Wgsl,
   isWgsl,
-} from '../types';
+} from '../types.js';
 import {
   getTypeForIndexAccess,
   getTypeForPropAccess,
   getTypeFromWgsl,
   numericLiteralToResource,
-} from './generationHelpers';
+} from './generationHelpers.js';
 
 const parenthesizedOps = [
   '==',
@@ -192,6 +193,24 @@ export function generateExpression(
         dataType: getTypeForPropAccess(target.dataType as Wgsl, property),
       };
     }
+
+    if (wgsl.isWgslArray(target.dataType)) {
+      if (property === 'length') {
+        if (target.dataType.elementCount === 0) {
+          // Dynamically-sized array
+          return {
+            value: `arrayLength(&${ctx.resolve(target.value)})`,
+            dataType: d.u32,
+          };
+        }
+
+        return {
+          value: String(target.dataType.elementCount),
+          dataType: abstractInt,
+        };
+      }
+    }
+
     // biome-ignore lint/suspicious/noExplicitAny: <sorry TypeScript>
     const propValue = (target.value as any)[property];
 

packages/typegpu/tests/array.test.ts

@@ -1,72 +1,74 @@
 import { BufferReader, BufferWriter } from 'typed-binary';
 import { describe, expect, it } from 'vitest';
-import { arrayOf, sizeOf, vec3f, vec3u } from '../src/data';
+import tgpu from '../src';
+import * as d from '../src/data';
 import { readData, writeData } from '../src/data/dataIO';
 import { StrictNameRegistry } from '../src/nameRegistry';
 import { resolve } from '../src/resolutionCtx';
 import type { Infer } from '../src/shared/repr';
+import { parse, parseResolved } from './utils/parseResolved';
 
 describe('array', () => {
   it('takes element alignment into account when measuring', () => {
-    const TestArray = arrayOf(vec3u, 3);
-    expect(sizeOf(TestArray)).toEqual(48);
+    const TestArray = d.arrayOf(d.vec3u, 3);
+    expect(d.sizeOf(TestArray)).toEqual(48);
   });
 
   it('aligns array elements when writing', () => {
-    const TestArray = arrayOf(vec3u, 3);
-    const buffer = new ArrayBuffer(sizeOf(TestArray));
+    const TestArray = d.arrayOf(d.vec3u, 3);
+    const buffer = new ArrayBuffer(d.sizeOf(TestArray));
     const writer = new BufferWriter(buffer);
 
     writeData(writer, TestArray, [
-      vec3u(1, 2, 3),
-      vec3u(4, 5, 6),
-      vec3u(7, 8, 9),
+      d.vec3u(1, 2, 3),
+      d.vec3u(4, 5, 6),
+      d.vec3u(7, 8, 9),
     ]);
     expect([...new Uint32Array(buffer)]).toEqual([
       1, 2, 3, 0, 4, 5, 6, 0, 7, 8, 9, 0,
     ]);
   });
 
   it('aligns array elements when reading', () => {
-    const TestArray = arrayOf(vec3u, 3);
-    const buffer = new ArrayBuffer(sizeOf(TestArray));
+    const TestArray = d.arrayOf(d.vec3u, 3);
+    const buffer = new ArrayBuffer(d.sizeOf(TestArray));
     const reader = new BufferReader(buffer);
 
     new Uint32Array(buffer).set([1, 2, 3, 0, 4, 5, 6, 0, 7, 8, 9, 0]);
 
     expect(readData(reader, TestArray)).toEqual([
-      vec3u(1, 2, 3),
-      vec3u(4, 5, 6),
-      vec3u(7, 8, 9),
+      d.vec3u(1, 2, 3),
+      d.vec3u(4, 5, 6),
+      d.vec3u(7, 8, 9),
     ]);
   });
 
   it('encodes and decodes arrays properly', () => {
-    const TestArray = arrayOf(vec3f, 5);
+    const TestArray = d.arrayOf(d.vec3f, 5);
 
-    const buffer = new ArrayBuffer(sizeOf(TestArray));
+    const buffer = new ArrayBuffer(d.sizeOf(TestArray));
 
     const value: Infer<typeof TestArray> = [
-      vec3f(1.5, 2, 3.5),
-      vec3f(),
-      vec3f(-1.5, 2, 3.5),
-      vec3f(1.5, -2, 3.5),
-      vec3f(1.5, 2, 15),
+      d.vec3f(1.5, 2, 3.5),
+      d.vec3f(),
+      d.vec3f(-1.5, 2, 3.5),
+      d.vec3f(1.5, -2, 3.5),
+      d.vec3f(1.5, 2, 15),
     ];
 
     writeData(new BufferWriter(buffer), TestArray, value);
     expect(readData(new BufferReader(buffer), TestArray)).toEqual(value);
   });
 
   it('throws when trying to read/write a runtime-sized array', () => {
-    const TestArray = arrayOf(vec3f, 0);
+    const TestArray = d.arrayOf(d.vec3f, 0);
 
-    expect(sizeOf(TestArray)).toBeNaN();
+    expect(d.sizeOf(TestArray)).toBeNaN();
 
     expect(() =>
       writeData(new BufferWriter(new ArrayBuffer(0)), TestArray, [
-        vec3f(),
-        vec3f(),
+        d.vec3f(),
+        d.vec3f(),
       ]),
     ).toThrow();
 
@@ -80,6 +82,70 @@ describe('array', () => {
   });
 
   it('throws when trying to nest runtime sized arrays', () => {
-    expect(() => arrayOf(arrayOf(vec3f, 0), 0)).toThrow();
+    expect(() => d.arrayOf(d.arrayOf(d.vec3f, 0), 0)).toThrow();
+  });
+});
+
+describe('array.length', () => {
+  it('works for dynamically-sized arrays in TGSL', () => {
+    const layout = tgpu.bindGroupLayout({
+      values: {
+        storage: (n: number) => d.arrayOf(d.f32, n),
+        access: 'mutable',
+      },
+    });
+
+    const foo = tgpu['~unstable'].fn([]).does(() => {
+      let acc = d.f32(1);
+      for (let i = 0; i < layout.bound.values.value.length; i++) {
+        layout.bound.values.value[i] = acc;
+        acc *= 2;
+      }
+    });
+
+    expect(parseResolved({ foo })).toEqual(
+      parse(/* wgsl */ `
+        @group(0) @binding(0) var <storage, read_write> values: array<f32>;
+
+        fn foo() {
+          var acc = f32(1);
+          for (var i = 0; (i < arrayLength(&values)); i++) {
+            values[i] = acc;
+            acc *= 2;
+          }
+        }
+      `),
+    );
+  });
+
+  it('works for statically-sized arrays in TGSL', () => {
+    const layout = tgpu.bindGroupLayout({
+      values: {
+        storage: d.arrayOf(d.f32, 128),
+        access: 'mutable',
+      },
+    });
+
+    const foo = tgpu['~unstable'].fn([]).does(() => {
+      let acc = d.f32(1);
+      for (let i = 0; i < layout.bound.values.value.length; i++) {
+        layout.bound.values.value[i] = acc;
+        acc *= 2;
+      }
+    });
+
+    expect(parseResolved({ foo })).toEqual(
+      parse(/* wgsl */ `
+        @group(0) @binding(0) var <storage, read_write> values: array<f32, 128>;
+
+        fn foo() {
+          var acc = f32(1);
+          for (var i = 0; (i < 128); i++) {
+            values[i] = acc;
+            acc *= 2;
+          }
+        }
+      `),
+    );
   });
 });