diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 2b5456e1dc..99547411f1 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -245,6 +245,9 @@ export class WgslGenerator implements ShaderGenerator { // used to detect `continue` and `break` nodes in loop body, as well as label // unrolled blocks with comments #unrollingChain: number[] = []; + // true while the innermost statement a `break` would exit is a `switch`, so + // that `break` survives unrolling the loop around the `switch` + #breakExitsSwitch = false; // prototype properties declare languageKey: string; @@ -1805,9 +1808,17 @@ ${this.ctx.pre}else ${alternate}`, const e = elements[i]; this.#unrollingChain = [...prevUnrollingChain, i]; - const resolvedBlock = this._blockStatement(blockified, { - [originalLoopVarName]: e, - }); + // A `break` here exits the unrolled loop, even inside a `switch`. + const prevBreakExitsSwitch = this.#breakExitsSwitch; + this.#breakExitsSwitch = false; + let resolvedBlock: ResolvedStatement; + try { + resolvedBlock = this._blockStatement(blockified, { + [originalLoopVarName]: e, + }); + } finally { + this.#breakExitsSwitch = prevBreakExitsSwitch; + } definesInNearestScope ||= resolvedBlock.definesInNearestScope; @@ -1933,10 +1944,13 @@ ${this.ctx.pre}else ${alternate}`, this.ctx.pushBlockScope(); this.ctx.indent(); this.ctx.indent(); + const prevBreakExitsSwitch = this.#breakExitsSwitch; + this.#breakExitsSwitch = true; try { const consequentStmts = consequent.map((s) => this._statement(s)); return [test, consequentStmts]; } finally { + this.#breakExitsSwitch = prevBreakExitsSwitch; this.ctx.dedent(); this.ctx.dedent(); this.ctx.popBlockScope(); @@ -2015,7 +2029,7 @@ ${stringifyNode(statement)}`); } if (statement[0] === NODE.break) { - if (this.#unrollingChain.length > 0) { + if (this.#unrollingChain.length > 0 && !this.#breakExitsSwitch) { throw new WgslTypeError('Cannot unroll loop containing `break`'); } return { diff --git a/packages/typegpu/tests/unroll.test.ts b/packages/typegpu/tests/unroll.test.ts index 7f5f4c3a42..a005014490 100644 --- a/packages/typegpu/tests/unroll.test.ts +++ b/packages/typegpu/tests/unroll.test.ts @@ -766,6 +766,128 @@ describe('tgpu.unroll', () => { `); }); + it('unrolls when `break` exits a `switch` inside the loop body', () => { + const f = (mode: number) => { + 'use gpu'; + let r = d.i32(0); + for (const foo of tgpu.unroll([1, 2])) { + switch (mode) { + case 0: + r += foo; + break; + default: + r -= foo; + break; + } + } + return r; + }; + + expect(tgpu.resolve([tgpu.fn([d.i32], d.i32)(f)])).toMatchInlineSnapshot(` + "fn f(mode: i32) -> i32 { + var r = 0i; + // unrolled iteration #0 + switch mode { + case 0i: { + r += 1i; + } + case default: { + r -= 1i; + } + } + // unrolled iteration #1 + switch mode { + case 0i: { + r += 2i; + } + case default: { + r -= 2i; + } + } + // --- + return r; + }" + `); + }); + + it('unrolls when `break` exits a comptime-folded `switch` inside the loop body', () => { + const f = () => { + 'use gpu'; + let r = d.i32(0); + for (const foo of tgpu.unroll([1, 2])) { + switch (foo) { + case 1: + r += 10; + break; + default: + r += foo; + } + } + return r; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + var r = 0i; + // unrolled iteration #0 + switch 1i { + case default: { + r += 10i; + } + } + // unrolled iteration #1 + switch 2i { + case default: { + r += 2i; + } + } + // --- + return r; + }" + `); + }); + + it('throws when `continue` is used inside a `switch` inside the loop body', () => { + const f = (mode: number) => { + 'use gpu'; + for (const foo of tgpu.unroll([1, 2])) { + switch (mode) { + case 0: + continue; + default: + break; + } + } + }; + + expect(() => tgpu.resolve([tgpu.fn([d.i32])(f)])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn:f: Cannot unroll loop containing \`continue\`] + `); + }); + + it('throws when `break` exits an unrolled loop nested in a `switch` case', () => { + const f = (mode: number) => { + 'use gpu'; + switch (mode) { + case 0: + for (const foo of tgpu.unroll([1, 2])) { + break; + } + break; + default: + break; + } + }; + + expect(() => tgpu.resolve([tgpu.fn([d.i32])(f)])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn:f: Cannot unroll loop containing \`break\`] + `); + }); + it('unrolling flag is set correctly', () => { const f = () => { 'use gpu';