feat: for ... of ... loop support (#1976)

cieplypolar · web-flow · commit 226e90a8c3f0 · 2026-02-11T17:02:54.000+01:00
diff --git a/apps/typegpu-docs/src/examples/simulation/boids/index.ts b/apps/typegpu-docs/src/examples/simulation/boids/index.ts
@@ -177,11 +177,7 @@ const simulate = (index: number) => {
   let alignmentCount = 0;
   let cohesionCount = 0;
 
-  for (let i = d.u32(0); i < layout.$.currentTrianglePos.length; i++) {
-    if (i === index) {
-      continue;
-    }
-    const other = layout.$.currentTrianglePos[i];
+  for (const other of layout.$.currentTrianglePos) {
     const dist = std.distance(instanceInfo.position, other.position);
     if (dist < params.$.separationDistance) {
       separation = std.add(
diff --git a/apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts b/apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts
@@ -107,9 +107,7 @@ const time = root.createUniform(d.f32);
 
 const isInsideObstacle = (x: number, y: number): boolean => {
   'use gpu';
-  for (let obsIdx = 0; obsIdx < MAX_OBSTACLES; obsIdx++) {
-    const obs = obstacles.$[obsIdx];
-
+  for (const obs of obstacles.$) {
     if (obs.enabled === 0) {
       continue;
     }
@@ -162,8 +160,7 @@ const computeVelocity = (x: number, y: number): d.v2f => {
   ];
   let dirChoiceCount = 1;
 
-  for (let i = 0; i < 4; i++) {
-    const offset = neighborOffsets[i];
+  for (const offset of neighborOffsets) {
     const neighborDensity = getCell(x + offset.x, y + offset.y);
     const cost = neighborDensity.z + d.f32(offset.y) * gravityCost;
 
diff --git a/packages/tinyest-for-wgsl/src/parsers.ts b/packages/tinyest-for-wgsl/src/parsers.ts
@@ -244,12 +244,7 @@ const Transpilers: Partial<
     }
 
     if (node.kind === 'const') {
-      if (init === undefined) {
-        throw new Error(
-          'Did not provide initial value in `const` declaration.',
-        );
-      }
-      return [NODE.const, id, init];
+      return init !== undefined ? [NODE.const, id, init] : [NODE.const, id];
     }
 
     return init !== undefined ? [NODE.let, id, init] : [NODE.let, id];
@@ -322,6 +317,13 @@ const Transpilers: Partial<
     return [NODE.while, condition, body];
   },
 
+  ForOfStatement(ctx, node) {
+    const loopVar = transpile(ctx, node.left) as tinyest.Const | tinyest.Let;
+    const iterable = transpile(ctx, node.right) as tinyest.Expression;
+    const body = transpile(ctx, node.body) as tinyest.Statement;
+    return [NODE.forOf, loopVar, iterable, body];
+  },
+
   ContinueStatement() {
     return [NODE.continue];
   },
diff --git a/packages/tinyest/src/nodes.ts b/packages/tinyest/src/nodes.ts
@@ -23,6 +23,7 @@ export const NodeTypeCatalog = {
   while: 15,
   continue: 16,
   break: 17,
+  forOf: 18,
 
   // rare
   arrayExpr: 100,
@@ -73,11 +74,13 @@ export type Let =
 /**
  * Represents a const statement
  */
-export type Const = readonly [
-  type: NodeTypeCatalog['const'],
-  identifier: string,
-  value: Expression,
-];
+export type Const =
+  | readonly [type: NodeTypeCatalog['const'], identifier: string]
+  | readonly [
+    type: NodeTypeCatalog['const'],
+    identifier: string,
+    value: Expression,
+  ];
 
 export type For = readonly [
   type: NodeTypeCatalog['for'],
@@ -97,6 +100,13 @@ export type Continue = readonly [type: NodeTypeCatalog['continue']];
 
 export type Break = readonly [type: NodeTypeCatalog['break']];
 
+export type ForOf = readonly [
+  type: NodeTypeCatalog['forOf'],
+  left: Const | Let,
+  right: Expression,
+  body: Statement,
+];
+
 /**
  * A union type of all statements
  */
@@ -110,7 +120,8 @@ export type Statement =
   | For
   | While
   | Continue
-  | Break;
+  | Break
+  | ForOf;
 
 //
 // Expression
diff --git a/packages/typegpu/src/data/compiledIO.ts b/packages/typegpu/src/data/compiledIO.ts
@@ -196,6 +196,10 @@ export function buildWriter(
   }
 
   if (wgsl.isVec(node)) {
+    if (wgsl.isVecBool(node)) {
+      throw new Error('Compiled writers do not support boolean vectors');
+    }
+
     const primitive = typeToPrimitive[node.type];
     let code = '';
     const writeFunc = primitiveToWriteFunction[primitive];
diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
@@ -1717,17 +1717,29 @@ export function isVec(
   | Vec2h
   | Vec2i
   | Vec2u
+  | Vec2b
   | Vec3f
   | Vec3h
   | Vec3i
   | Vec3u
+  | Vec3b
   | Vec4f
   | Vec4h
   | Vec4i
-  | Vec4u {
+  | Vec4u
+  | Vec4b {
   return isVec2(value) || isVec3(value) || isVec4(value);
 }
 
+export function isVecBool(
+  value: unknown,
+): value is
+  | Vec2b
+  | Vec3b
+  | Vec4b {
+  return isVec(value) && value.type.includes('b');
+}
+
 export function isMatInstance(value: unknown): value is AnyMatInstance {
   const v = value as AnyMatInstance | undefined;
   return isMarkedInternal(v) &&
diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts
@@ -41,6 +41,7 @@ import type { ShaderGenerator } from './shaderGenerator.ts';
 import { createPtrFromOrigin, implicitFrom, ptrFn } from '../data/ptr.ts';
 import { RefOperator } from '../data/ref.ts';
 import { constant } from '../core/constant/tgpuConstant.ts';
+import { arrayLength } from '../std/array.ts';
 import { AutoStruct } from '../data/autoStruct.ts';
 import { mathToStd } from './math.ts';
 
@@ -1137,6 +1138,131 @@ ${this.ctx.pre}else ${alternate}`;
       return `${this.ctx.pre}while (${conditionStr}) ${bodyStr}`;
     }
 
+    if (statement[0] === NODE.forOf) {
+      const [_, loopVar, iterable, body] = statement;
+      const iterableSnippet = this.expression(iterable);
+
+      if (isEphemeralSnippet(iterableSnippet)) {
+        throw new Error(
+          '`for ... of ...` loops only support iterables stored in variables',
+        );
+      }
+
+      // Our index name will be some element from infinite sequence (i, ii, iii, ...).
+      // If user defines `i` and `ii` before `for ... of ...` loop, then our index name will be `iii`.
+      // If user defines `i` inside `for ... of ...` then it will be scoped to a new block,
+      // so we can safely use `i`.
+      let index = 'i'; // it will be valid name, no need to call this.ctx.makeNameValid
+      while (this.ctx.getById(index) !== null) {
+        index += 'i';
+      }
+
+      const elementSnippet = accessIndex(
+        iterableSnippet,
+        snip(index, u32, 'runtime'),
+      );
+      if (!elementSnippet) {
+        throw new WgslTypeError(
+          '`for ... of ...` loops only support array or vector iterables',
+        );
+      }
+
+      const iterableDataType = iterableSnippet.dataType;
+      let elementCountSnippet: Snippet;
+      let elementType = elementSnippet.dataType;
+
+      if (elementType === UnknownData) {
+        throw new WgslTypeError(
+          stitch`The elements in iterable ${iterableSnippet} are of unknown type`,
+        );
+      }
+
+      if (wgsl.isWgslArray(iterableDataType)) {
+        elementCountSnippet = iterableDataType.elementCount > 0
+          ? snip(
+            `${iterableDataType.elementCount}`,
+            u32,
+            'constant',
+          )
+          : arrayLength[$gpuCallable].call(this.ctx, [iterableSnippet]);
+      } else if (wgsl.isVec(iterableDataType)) {
+        elementCountSnippet = snip(
+          `${Number(iterableDataType.type.match(/\d/))}`,
+          u32,
+          'constant',
+        );
+      } else {
+        throw new WgslTypeError(
+          '`for ... of ...` loops only support array or vector iterables',
+        );
+      }
+
+      if (loopVar[0] !== NODE.const) {
+        throw new WgslTypeError(
+          'Only `for (const ... of ... )` loops are supported',
+        );
+      }
+
+      // If it's ephemeral, it's a value that cannot change. If it's a reference, we take
+      // an implicit pointer to it
+      let loopVarKind = 'let';
+      const loopVarName = this.ctx.makeNameValid(loopVar[1]);
+
+      if (!isEphemeralSnippet(elementSnippet)) {
+        if (elementSnippet.origin === 'constant-tgpu-const-ref') {
+          loopVarKind = 'const';
+        } else if (elementSnippet.origin === 'runtime-tgpu-const-ref') {
+          loopVarKind = 'let';
+        } else {
+          loopVarKind = 'let';
+          if (!wgsl.isPtr(elementType)) {
+            const ptrType = createPtrFromOrigin(
+              elementSnippet.origin,
+              concretize(elementType as wgsl.AnyWgslData) as wgsl.StorableData,
+            );
+            invariant(
+              ptrType !== undefined,
+              `Creating pointer type from origin ${elementSnippet.origin}`,
+            );
+            elementType = ptrType;
+          }
+
+          elementType = implicitFrom(elementType as wgsl.Ptr);
+        }
+      }
+
+      const loopVarSnippet = snip(
+        loopVarName,
+        elementType,
+        elementSnippet.origin,
+      );
+      this.ctx.defineVariable(loopVarName, loopVarSnippet);
+
+      const forStr = stitch`${this.ctx.pre}for (var ${index} = 0u; ${index} < ${
+        tryConvertSnippet(this.ctx, elementCountSnippet, u32, false)
+      }; ${index}++) {`;
+
+      this.ctx.indent();
+
+      const loopVarDeclStr =
+        stitch`${this.ctx.pre}${loopVarKind} ${loopVarName} = ${
+          tryConvertSnippet(
+            this.ctx,
+            elementSnippet,
+            elementType,
+            false,
+          )
+        };`;
+
+      const bodyStr = `${this.ctx.pre}${
+        this.block(blockifySingleStatement(body))
+      }`;
+
+      this.ctx.dedent();
+
+      return stitch`${forStr}\n${loopVarDeclStr}\n${bodyStr}\n${this.ctx.pre}}`;
+    }
+
     if (statement[0] === NODE.continue) {
       return `${this.ctx.pre}continue;`;
     }
diff --git a/packages/typegpu/tests/examples/individual/boids.test.ts b/packages/typegpu/tests/examples/individual/boids.test.ts
@@ -46,22 +46,21 @@ describe('boids example', () => {
         var cohesion = vec2f();
         var alignmentCount = 0;
         var cohesionCount = 0;
-        for (var i = 0u; (i < arrayLength(&currentTrianglePos)); i++) {
-          if ((i == index)) {
-            continue;
-          }
+        for (var i = 0u; i < arrayLength((&currentTrianglePos)); i++) {
           let other = (&currentTrianglePos[i]);
-          let dist = distance(instanceInfo.position, (*other).position);
-          if ((dist < paramsBuffer.separationDistance)) {
-            separation = (separation + (instanceInfo.position - (*other).position));
-          }
-          if ((dist < paramsBuffer.alignmentDistance)) {
-            alignment = (alignment + (*other).velocity);
-            alignmentCount++;
-          }
-          if ((dist < paramsBuffer.cohesionDistance)) {
-            cohesion = (cohesion + (*other).position);
-            cohesionCount++;
+          {
+            let dist = distance(instanceInfo.position, (*other).position);
+            if ((dist < paramsBuffer.separationDistance)) {
+              separation = (separation + (instanceInfo.position - (*other).position));
+            }
+            if ((dist < paramsBuffer.alignmentDistance)) {
+              alignment = (alignment + (*other).velocity);
+              alignmentCount++;
+            }
+            if ((dist < paramsBuffer.cohesionDistance)) {
+              cohesion = (cohesion + (*other).position);
+              cohesionCount++;
+            }
           }
         }
         if ((alignmentCount > 0i)) {
diff --git a/packages/typegpu/tests/examples/individual/fluid-double-buffering.test.ts b/packages/typegpu/tests/examples/individual/fluid-double-buffering.test.ts
diff --git a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts
