fix: Better std.range internal variables type inference (#2384)

Author: cieplypolar

packages/typegpu/src/tgsl/forOfUtils.ts

@@ -63,11 +63,14 @@ export function getRangeSnippets(
   const { value, dataType } = iterableSnippet;
 
   if (isTgpuRange(value)) {
+    const { start, end, step } = value;
+    const dataType = [start, end, step].every((v) => v >= 0) ? u32 : i32;
+
     return {
-      start: snip(value.start, i32, 'constant'),
-      end: snip(value.end, i32, 'constant'),
-      step: snip(value.step, i32, 'constant'),
-      comparison: value.step < 0 ? '>' : '<',
+      start: snip(start, dataType, 'constant'),
+      end: snip(end, dataType, 'constant'),
+      step: snip(step, dataType, 'constant'),
+      comparison: step < 0 ? '>' : '<',
     };
   }
 

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -1263,7 +1263,7 @@ ${this.ctx.pre}else ${alternate}`;
 
         if (isTgpuRange(iterableSnippet.value)) {
           bodyStr = this._block(blockified, {
-            [originalLoopVarName]: snip(index, u32, 'runtime'),
+            [originalLoopVarName]: snip(index, range.start.dataType, 'runtime'), // range.start, .end , .step have the same dataType
           });
         } else {
           this.ctx.indent();

packages/typegpu/tests/std/range.test.ts

@@ -107,7 +107,7 @@ describe('on the GPU', () => {
     expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
       "fn main() -> f32 {
         var result = 0f;
-        for (var i = 0i; i < 8i; i += 2i) {
+        for (var i = 0u; i < 8u; i += 2u) {
           result += f32(i);
         }
         for (var i = 10i; i > -10i; i += -1i) {

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

@@ -979,6 +979,30 @@ describe('wgslGenerator', () => {
     `);
   });
 
+  it('handles "for ... of ..." over `std.range`', ({ root }) => {
+    const testFn = tgpu.fn(
+      [],
+      d.f32,
+    )(() => {
+      'use gpu';
+      for (const bounce of std.range(12)) {
+        const test = d.u32(2) + bounce;
+        return d.f32(test);
+      }
+      return 0;
+    });
+
+    expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(`
+      "fn testFn() -> f32 {
+        for (var i = 0u; i < 12u; i += 1u) {
+          let test = (2u + i);
+          return f32(test);
+        }
+        return 0f;
+      }"
+    `);
+  });
+
   it('creates correct resources for lazy values and slots', () => {
     const testFn = tgpu.fn([], d.vec4u)(() => lazyV4u.$);