diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 0b6e1aa8dc..a972253caf 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -1728,22 +1728,20 @@ ${this.ctx.pre}else ${alternate}`, if (!isKnownAtComptime(range.end)) { throw new Error('Cannot unroll loop. Length of iterable is unknown at comptime.'); } - - const length = range.end.value as number; - if (length === 0) { - return { code: '', definesInNearestScope: false }; - } - const { value } = iterableSnippet; const elements = isTgpuRange(value) ? value.map((i) => coerceToSnippet(i)) : value instanceof ArrayExpression ? value.elements - : Array.from({ length }, (_, i) => + : Array.from({ length: range.end.value as number }, (_, i) => forOfUtils.getElementSnippet(iterableSnippet, snip(i, u32, 'constant')), ); + if (elements.length === 0) { + return { code: '', definesInNearestScope: false }; + } + const firstElement = elements[0] as Snippet; if (!isAlias(firstElement) && !wgsl.isNaturallyEphemeral(firstElement.dataType)) { throw new WgslTypeError( diff --git a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts index 3b75915f89..29d28b740d 100644 --- a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts +++ b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts @@ -2347,4 +2347,49 @@ describe('WgslGenerator', () => { `); }); }); + + it('unrolls a descending range ending at zero', () => { + const main = () => { + 'use gpu'; + let sum = 0; + for (const i of tgpu.unroll(std.range(3, 0, -1))) { + sum += i; + } + return sum; + }; + + expect(main()).toBe(6); + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "fn main() -> i32 { + var sum = 0; + // unrolled iteration #0 + sum += 3i; + // unrolled iteration #1 + sum += 2i; + // unrolled iteration #2 + sum += 1i; + // --- + return sum; + }" + `); + }); + + it('unrolls an empty range with a nonzero endpoint', () => { + const main = () => { + 'use gpu'; + let sum = 0; + for (const i of tgpu.unroll(std.range(3, 3))) { + sum += i; + } + return sum; + }; + + expect(main()).toBe(0); + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "fn main() -> i32 { + let sum = 0; + return sum; + }" + `); + }); });