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

apps/typegpu-docs/src/examples/rendering/point-light-shadow/index.ts

@@ -182,8 +182,8 @@ const fragmentMain = tgpu['~unstable'].fragmentFn({
   const diskRadius = shadowParams.$.diskRadius;
 
   let visibilityAcc = 0.0;
-  for (let i = 0; i < PCF_SAMPLES; i++) {
-    const o = samplesUniform.$[i].xy.mul(diskRadius);
+  for (const sample of samplesUniform.$) {
+    const o = sample.xy.mul(diskRadius);
 
     const sampleDir = dir
       .add(right.mul(o.x))

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/tgsl/wgslGenerator.ts

@@ -1043,6 +1043,108 @@ ${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',
+        );
+      }
+
+      // if user defines a variable named 'i', it will be scoped to a new block,
+      // shadowing our 'i'
+      const index = this.ctx.makeNameValid('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 elementCount: number;
+      let elementType = elementSnippet?.dataType;
+      if (wgsl.isWgslArray(iterableDataType)) {
+        elementCount = iterableDataType.elementCount;
+      } else if (wgsl.isVec(iterableDataType)) {
+        elementCount = Number(iterableDataType.type.match(/\d/));
+      } 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 iterableStr = this.ctx.resolve(
+        iterableSnippet.value,
+        iterableDataType,
+      ).value;
+
+      const forStr =
+        `${this.ctx.pre}for (var ${index} = 0; ${index} < ${elementCount}; ${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;`;
     }