diff --git a/packages/typegpu/src/core/rawCodeSnippet/tgpuRawCodeSnippet.ts b/packages/typegpu/src/core/rawCodeSnippet/tgpuRawCodeSnippet.ts index 30a105e90c..7995f86f67 100644 --- a/packages/typegpu/src/core/rawCodeSnippet/tgpuRawCodeSnippet.ts +++ b/packages/typegpu/src/core/rawCodeSnippet/tgpuRawCodeSnippet.ts @@ -112,13 +112,25 @@ class TgpuRawCodeSnippetImpl implements TgpuRawCodeS this.#expression, ); - return snip(replacedExpression, this.dataType, this.origin, this.possibleSideEffects); + return snip( + replacedExpression, + this.dataType, + this.origin, + this.possibleSideEffects, + /* mayDeclareVariables */ true, + ); }, }), { codegenMode: { getBaseSnippet(trackingProxy) { - return snip(trackingProxy, this.dataType, this.origin, this.possibleSideEffects); + return snip( + trackingProxy, + this.dataType, + this.origin, + this.possibleSideEffects, + /* mayDeclareVariables */ true, + ); }, }, normalMode: { diff --git a/packages/typegpu/src/data/snippet.ts b/packages/typegpu/src/data/snippet.ts index 35db60b5f0..a8366e1650 100644 --- a/packages/typegpu/src/data/snippet.ts +++ b/packages/typegpu/src/data/snippet.ts @@ -103,6 +103,11 @@ export interface Snippet { * that mutates memory, or synchronizes threads, has side effects. */ readonly possibleSideEffects: boolean; + /** + * Whether opaque source code may introduce names in the surrounding scope. + * Statement blocks containing such snippets must not be inlined. + */ + readonly mayDeclareVariables?: true; } export interface ResolvedSnippet extends Snippet { @@ -123,11 +128,15 @@ class SnippetImpl implements Snippet { dataType: BaseData | UnknownData, origin: Origin, possibleSideEffects: boolean, + mayDeclareVariables: boolean, ) { this.value = value; this.dataType = dataType; this.origin = origin; this.possibleSideEffects = possibleSideEffects; + if (mayDeclareVariables) { + Object.defineProperty(this, 'mayDeclareVariables', { value: true }); + } } } @@ -144,18 +153,21 @@ export function snip( dataType: BaseData, origin: Origin, possibleSideEffects?: boolean, + mayDeclareVariables?: boolean, ): ResolvedSnippet; export function snip( value: unknown, dataType: BaseData | UnknownData, origin: Origin, possibleSideEffects?: boolean, + mayDeclareVariables?: boolean, ): Snippet; export function snip( value: unknown, dataType: BaseData | UnknownData, origin: Origin, possibleSideEffects: boolean = true, + mayDeclareVariables: boolean = false, ): Snippet | ResolvedSnippet { if (DEV && isSnippet(value)) { // An early error, but not worth checking every time in production @@ -168,6 +180,7 @@ export function snip( undecorate(dataType as BaseData), origin, possibleSideEffects, + mayDeclareVariables, ); } @@ -177,13 +190,25 @@ export function withDataType( ): ResolvedSnippet; export function withDataType(dataType: BaseData | UnknownData, snippet: Snippet): Snippet; export function withDataType(dataType: BaseData | UnknownData, snippet: Snippet): Snippet { - return new SnippetImpl(snippet.value, dataType, snippet.origin, snippet.possibleSideEffects); + return new SnippetImpl( + snippet.value, + dataType, + snippet.origin, + snippet.possibleSideEffects, + snippet.mayDeclareVariables ?? false, + ); } export function withValue(value: string, snippet: Snippet): ResolvedSnippet; export function withValue(value: unknown, snippet: Snippet): Snippet; export function withValue(value: unknown, snippet: Snippet): Snippet { - return new SnippetImpl(value, snippet.dataType, snippet.origin, snippet.possibleSideEffects); + return new SnippetImpl( + value, + snippet.dataType, + snippet.origin, + snippet.possibleSideEffects, + snippet.mayDeclareVariables ?? false, + ); } export function withSideEffects( @@ -192,7 +217,13 @@ export function withSideEffects( ): ResolvedSnippet; export function withSideEffects(possibleSideEffects: boolean, snippet: Snippet): Snippet; export function withSideEffects(possibleSideEffects: boolean, snippet: Snippet): Snippet { - return new SnippetImpl(snippet.value, snippet.dataType, snippet.origin, possibleSideEffects); + return new SnippetImpl( + snippet.value, + snippet.dataType, + snippet.origin, + possibleSideEffects, + snippet.mayDeclareVariables ?? false, + ); } export function noSideEffects(snippet: ResolvedSnippet): ResolvedSnippet; diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 33d64ec786..d9bce7ff5c 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -1833,7 +1833,10 @@ ${this.ctx.pre}else ${alternate}`, const expr = this._expression(statement); const resolved = expr.value !== undefined && expr.value !== null ? this.ctx.resolveSnippet(expr).value : ''; - return { code: resolved ? `${this.ctx.pre}${resolved};` : '', definesInNearestScope: false }; + return { + code: resolved ? `${this.ctx.pre}${resolved};` : '', + definesInNearestScope: expr.mayDeclareVariables ?? false, + }; } /** diff --git a/packages/typegpu/tests/unroll.test.ts b/packages/typegpu/tests/unroll.test.ts index 7a9898f1d5..3da360e239 100644 --- a/packages/typegpu/tests/unroll.test.ts +++ b/packages/typegpu/tests/unroll.test.ts @@ -927,4 +927,32 @@ describe('tgpu.unroll', () => { }" `); }); + it('keeps scopes around unrolled raw code snippets', () => { + const snippet = tgpu['~unstable'].rawCodeSnippet('const value = 1', d.Void); + + const main = () => { + 'use gpu'; + for (const _ of tgpu.unroll([1, 2, 3])) { + snippet.$; + } + }; + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "fn main() { + // unrolled iteration #0 + { + const value = 1; + } + // unrolled iteration #1 + { + const value = 1; + } + // unrolled iteration #2 + { + const value = 1; + } + // --- + }" + `); + }); });