feat: for ... of ... loop support

apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts

@@ -103,9 +103,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;
     }
@@ -158,8 +156,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;
 

packages/tinyest-for-wgsl/src/parsers.ts

@@ -236,12 +236,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];
@@ -314,6 +309,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];
   },

packages/tinyest/src/nodes.ts

@@ -23,6 +23,7 @@ export const NodeTypeCatalog = {
   while: 15,
   continue: 16,
   break: 17,
+  forOf: 18,
 
   // rare
   arrayExpr: 100,
@@ -72,11 +73,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'],
@@ -96,6 +99,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
  */
@@ -109,7 +119,8 @@ export type Statement =
   | For
   | While
   | Continue
-  | Break;
+  | Break
+  | ForOf;
 
 //
 // Expression

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];

packages/typegpu/src/data/wgslTypes.ts

@@ -1688,17 +1688,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) &&

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -21,9 +21,8 @@ import {
 import * as wgsl from '../data/wgslTypes.ts';
 import { invariant, ResolutionError, WgslTypeError } from '../errors.ts';
 import { getName } from '../shared/meta.ts';
-import { isMarkedInternal } from '../shared/symbols.ts';
+import { $internal, isMarkedInternal } from '../shared/symbols.ts';
 import { safeStringify } from '../shared/stringify.ts';
-import { $internal } from '../shared/symbols.ts';
 import { pow } from '../std/numeric.ts';
 import { add, div, mul, neg, sub } from '../std/operators.ts';
 import { type FnArgsConversionHint, isKnownAtComptime } from '../types.ts';
@@ -44,6 +43,7 @@ import type { DualFn } from '../data/dualFn.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';
 
 const { NodeTypeCatalog: NODE } = tinyest;
 
@@ -1058,6 +1058,119 @@ ${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 (wgsl.isWgslArray(iterableDataType)) {
+        elementCountSnippet = iterableDataType.elementCount > 0
+          ? snip(
+            `${iterableDataType.elementCount}`,
+            u32,
+            'constant',
+          )
+          : arrayLength[$internal].gpuImpl(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);
+        }
+      }
+
+      const loopVarSnippet = snip(
+        loopVarName,
+        elementType,
+        elementSnippet.origin,
+      );
+      this.ctx.defineVariable(loopVarName, loopVarSnippet);
+
+      const forStr = stitch`${this.ctx.pre}for (var ${index} = 0; ${index} < ${
+        tryConvertSnippet(elementCountSnippet, u32, false)
+      }; ${index}++) {`;
+
+      this.ctx.indent();
+
+      const loopVarDeclStr =
+        stitch`${this.ctx.pre}${loopVarKind} ${loopVarName} = ${
+          tryConvertSnippet(elementSnippet, elementType as AnyData, 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;`;
     }