cleanup

Author: cieplypolar

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -1054,12 +1054,11 @@ ${this.ctx.pre}else ${alternate}`;
       }
 
       const iterableDataType = iterableSnippet.dataType;
-
       let elementCount: number;
-      let elementType: wgsl.BaseData;
+      let elementType: wgsl.AnyWgslData;
       if (wgsl.isWgslArray(iterableDataType)) {
         elementCount = iterableDataType.elementCount;
-        elementType = iterableDataType.elementType;
+        elementType = iterableDataType.elementType as wgsl.AnyWgslData;
       } else if (wgsl.isVec(iterableDataType)) {
         elementType = iterableDataType.primitive;
         elementCount = Number(iterableDataType.type.match(/\d/));
@@ -1068,29 +1067,28 @@ ${this.ctx.pre}else ${alternate}`;
           '`for ... of ...` loops only support array or vector iterables',
         );
       }
-      const loopVarKind = loopVar[0] === tinyest.NodeTypeCatalog.const &&
-          wgsl.isNaturallyEphemeral(elementType)
-        ? 'let'
-        : 'var';
 
       /*
        * 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 loopVarKind = loopVar[0] === tinyest.NodeTypeCatalog.const &&
+          wgsl.isNaturallyEphemeral(elementType)
+        ? 'let'
+        : 'var';
       const loopVarName = this.ctx.makeNameValid(loopVar[1]);
       const loopVarSnippet = snip(
         loopVarName,
-        elementType as AnyData,
+        elementType,
         'runtime',
       );
       this.ctx.defineVariable(loopVarName, loopVarSnippet);
 
       const iterableStr = this.ctx.resolve(
         iterableSnippet.value,
-        // it's vector or wgslArray
-        iterableSnippet.dataType as AnyData,
+        iterableDataType,
       ).value;
 
       const forStr =

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

@@ -449,7 +449,38 @@ describe('wgslGenerator', () => {
     `);
   });
 
-  it('creates correct code for for ... of ... statement with array of primitives', () => {
+  it('parses correctly "for ... of ..." statements', () => {
+    const main1 = () => {
+      'use gpu';
+      const arr = [1, 2, 3];
+      for (const foo of arr) {
+        // biome-ignore lint/complexity/noUselessContinue: it's a part of the test
+        continue;
+      }
+    };
+
+    const main2 = () => {
+      'use gpu';
+      const arr = [1, 2, 3];
+      // biome-ignore lint/style/useConst: it's a part of the test
+      for (let foo of arr) {
+        // biome-ignore lint/complexity/noUselessContinue: it's a part of the test
+        continue;
+      }
+    };
+
+    const parsed1 = getMetaData(main1)?.ast?.body;
+    expect(JSON.stringify(parsed1)).toMatchInlineSnapshot(
+      `"[0,[[13,"arr",[100,[[5,"1"],[5,"2"],[5,"3"]]]],[18,[13,"foo"],"arr",[0,[[16]]]]]]"`,
+    );
+
+    const parsed2 = getMetaData(main2)?.ast?.body;
+    expect(JSON.stringify(parsed2)).toMatchInlineSnapshot(
+      `"[0,[[13,"arr",[100,[[5,"1"],[5,"2"],[5,"3"]]]],[18,[12,"foo"],"arr",[0,[[16]]]]]]"`,
+    );
+  });
+
+  it('creates correct code for "for ... of ..." statement using array of primitives', () => {
     const main = () => {
       'use gpu';
       const arr = d.arrayOf(d.f32, 3)([1, 2, 3]);
@@ -460,28 +491,54 @@ describe('wgslGenerator', () => {
       }
     };
 
+    expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
+      "fn main() {
+        var arr = array<f32, 3>(1f, 2f, 3f);
+        var res = 0f;
+        for (var i = 0; i < 3; i++) {
+          let foo = arr[i];
+          {
+            res += foo;
+            let i = res;
+          }
+        }
+      }"
+    `);
+  });
+
+  it('creates correct code for "for ... of ..." statement using array of non-primitives', () => {
+    const main = () => {
+      'use gpu';
+      const arr = d.arrayOf(d.vec2f, 3)([d.vec2f(1), d.vec2f(2), d.vec2f(3)]);
+      let res = 0;
+      for (const foo of arr) {
+        res += foo.x;
+        const i = res;
+      }
+    };
+
     const parsed = getMetaData(main)?.ast?.body;
 
     expect(JSON.stringify(parsed)).toMatchInlineSnapshot(
-      `"[0,[[13,"arr",[6,[6,[7,"d","arrayOf"],[[7,"d","f32"],[5,"3"]]],[[100,[[5,"1"],[5,"2"],[5,"3"]]]]]],[12,"res",[6,[7,"d","f32"],[]]],[18,[13,"foo"],"arr",[0,[[2,"res","+=","foo"],[13,"i","res"]]]]]]"`,
+      `"[0,[[13,"arr",[6,[6,[7,"d","arrayOf"],[[7,"d","vec2f"],[5,"3"]]],[[100,[[6,[7,"d","vec2f"],[[5,"1"]]],[6,[7,"d","vec2f"],[[5,"2"]]],[6,[7,"d","vec2f"],[[5,"3"]]]]]]]],[12,"res",[5,"0"]],[18,[13,"foo"],"arr",[0,[[2,"res","+=",[7,"foo","x"]],[13,"i","res"]]]]]]"`,
     );
 
     expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
       "fn main() {
-        var arr = array<f32, 3>(1f, 2f, 3f);
-        var res = 0f;
+        var arr = array<vec2f, 3>(vec2f(1), vec2f(2), vec2f(3));
+        var res = 0;
         for (var i = 0; i < 3; i++) {
-          let foo = arr[i];
+          var foo = arr[i];
           {
-            res += foo;
+            res += i32(foo.x);
             let i = res;
           }
         }
       }"
     `);
   });
 
-  it('creates correct code for for ... of ... statement derived, comptime iterables', () => {
+  it('creates correct code for "for ... of ..." statement using derived and comptime iterables', () => {
     const comptimeVec = tgpu['~unstable'].comptime(() => d.vec4f(1, 8, 8, 2));
 
     const main = () => {
@@ -519,39 +576,7 @@ describe('wgslGenerator', () => {
     `);
   });
 
-  it('creates correct code for for ... of ... statement with array of non-primitives', () => {
-    const main = () => {
-      'use gpu';
-      const arr = d.arrayOf(d.vec2f, 3)([d.vec2f(1), d.vec2f(2), d.vec2f(3)]);
-      let res = 0;
-      for (const foo of arr) {
-        res += foo.x;
-        const i = res;
-      }
-    };
-
-    const parsed = getMetaData(main)?.ast?.body;
-
-    expect(JSON.stringify(parsed)).toMatchInlineSnapshot(
-      `"[0,[[13,"arr",[6,[6,[7,"d","arrayOf"],[[7,"d","vec2f"],[5,"3"]]],[[100,[[6,[7,"d","vec2f"],[[5,"1"]]],[6,[7,"d","vec2f"],[[5,"2"]]],[6,[7,"d","vec2f"],[[5,"3"]]]]]]]],[12,"res",[5,"0"]],[18,[13,"foo"],"arr",[0,[[2,"res","+=",[7,"foo","x"]],[13,"i","res"]]]]]]"`,
-    );
-
-    expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
-      "fn main() {
-        var arr = array<vec2f, 3>(vec2f(1), vec2f(2), vec2f(3));
-        var res = 0;
-        for (var i = 0; i < 3; i++) {
-          var foo = arr[i];
-          {
-            res += i32(foo.x);
-            let i = res;
-          }
-        }
-      }"
-    `);
-  });
-
-  it('creates correct code for for ... of ... statement with vector iterables', () => {
+  it('creates correct code for "for ... of ..." statement using vector iterables', () => {
     const main = () => {
       'use gpu';
       const v1 = d.vec4f(1, 2, 3, 4);
@@ -605,7 +630,7 @@ describe('wgslGenerator', () => {
     `);
   });
 
-  it('creates correct code for for ... of ... statement with struct member iterable', () => {
+  it('creates correct code for "for ... of ..." statement using a struct member iterable', () => {
     const TestStruct = d.struct({
       arr: d.arrayOf(d.f32, 4),
     });
@@ -636,7 +661,7 @@ describe('wgslGenerator', () => {
     `);
   });
 
-  it('throws error when for ... of ... iterable is ephemeral', () => {
+  it('throws error when "for ... of ..." statement uses an ephemeral iterable', () => {
     const main = () => {
       'use gpu';
       for (const foo of [1, 2, 3]) {
@@ -653,7 +678,7 @@ describe('wgslGenerator', () => {
     `);
   });
 
-  it('throws error when for ... of ... iterable is not an array or a vector', () => {
+  it('throws error when "for ... of ..." statement uses iterable that is not an array or a vector', () => {
     const TestStruct = d.struct({
       x: d.u32,
       y: d.f32,