diff --git a/apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts b/apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts index fb47dade46..ce72e5caf7 100644 --- a/apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts +++ b/apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts @@ -454,7 +454,6 @@ function makePipelines( outputGridMutable: TgpuBufferMutable, ) { const initWorldPipeline = root['~unstable'] - .with(inputGridSlot, outputGridMutable) .with(outputGridSlot, outputGridMutable) .createGuardedComputePipeline((xu, yu) => { 'use gpu'; diff --git a/packages/typegpu/src/core/function/tgpuFn.ts b/packages/typegpu/src/core/function/tgpuFn.ts index 636efeaf46..ff6ec0e097 100644 --- a/packages/typegpu/src/core/function/tgpuFn.ts +++ b/packages/typegpu/src/core/function/tgpuFn.ts @@ -291,7 +291,7 @@ function createBoundFunction( slot: TgpuSlot | TgpuAccessor, value: unknown, ): TgpuFn { - return createBoundFunction(fn, [ + return createBoundFunction(innerFn, [ ...pairs, [isAccessor(slot) ? slot.slot : slot, value], ]); diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index 4b13aa8afb..18a7815828 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -76,6 +76,7 @@ import type { ItemLayer, ItemStateStack, ResolutionCtx, + StackLayer, TgpuShaderStage, Wgsl, } from './types.ts'; @@ -103,23 +104,8 @@ export type ResolutionCtxImplOptions = { readonly namespace: Namespace; }; -type SlotBindingLayer = { - type: 'slotBinding'; - bindingMap: WeakMap, unknown>; -}; - -type BlockScopeLayer = { - type: 'blockScope'; - declarations: Map; -}; - class ItemStateStackImpl implements ItemStateStack { - private _stack: ( - | ItemLayer - | SlotBindingLayer - | FunctionScopeLayer - | BlockScopeLayer - )[] = []; + private _stack: StackLayer[] = []; private _itemDepth = 0; get itemDepth(): number { @@ -146,10 +132,6 @@ class ItemStateStackImpl implements ItemStateStack { }); } - popItem() { - this.pop('item'); - } - pushSlotBindings(pairs: SlotValuePair[]) { this._stack.push({ type: 'slotBinding', @@ -157,10 +139,6 @@ class ItemStateStackImpl implements ItemStateStack { }); } - popSlotBindings() { - this.pop('slotBinding'); - } - pushFunctionScope( functionType: 'normal' | TgpuShaderStage, args: Snippet[], @@ -182,10 +160,6 @@ class ItemStateStackImpl implements ItemStateStack { return scope; } - popFunctionScope() { - this.pop('functionScope'); - } - pushBlockScope() { this._stack.push({ type: 'blockScope', @@ -193,20 +167,19 @@ class ItemStateStackImpl implements ItemStateStack { }); } - popBlockScope() { - this.pop('blockScope'); - } - - pop(type?: (typeof this._stack)[number]['type']) { + pop(type: T): Extract; + pop(): StackLayer | undefined; + pop(type?: StackLayer['type']) { const layer = this._stack[this._stack.length - 1]; if (!layer || (type && layer.type !== type)) { throw new Error(`Internal error, expected a ${type} layer to be on top.`); } - this._stack.pop(); + const poppedValue = this._stack.pop(); if (type === 'item') { this._itemDepth--; } + return poppedValue; } readSlot(slot: TgpuSlot): T | undefined { @@ -452,7 +425,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { } popBlockScope() { - this._itemStateStack.popBlockScope(); + this._itemStateStack.pop('blockScope'); } generateLog(op: string, args: Snippet[]): Snippet { @@ -560,7 +533,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { }; } finally { if (fnScopePushed) { - this._itemStateStack.popFunctionScope(); + this._itemStateStack.pop('functionScope'); } this.#namespaceInternal.nameRegistry.popFunctionScope(); } @@ -612,7 +585,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { try { return callback(); } finally { - this._itemStateStack.popSlotBindings(); + this._itemStateStack.pop('slotBinding'); } } @@ -700,7 +673,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { throw new ResolutionError(err, [derived]); } finally { - this._itemStateStack.popItem(); + this._itemStateStack.pop('item'); } } @@ -772,7 +745,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { throw new ResolutionError(err, [item]); } finally { - this._itemStateStack.popItem(); + this._itemStateStack.pop('item'); } } diff --git a/packages/typegpu/src/tgsl/generationHelpers.ts b/packages/typegpu/src/tgsl/generationHelpers.ts index 3b5061e24c..8dd19df4bf 100644 --- a/packages/typegpu/src/tgsl/generationHelpers.ts +++ b/packages/typegpu/src/tgsl/generationHelpers.ts @@ -1,12 +1,20 @@ +import { $internal, $resolve } from '../../src/shared/symbols.ts'; import { type AnyData, UnknownData } from '../data/dataTypes.ts'; import { abstractFloat, abstractInt, bool, f32, i32 } from '../data/numeric.ts'; import { isRef } from '../data/ref.ts'; -import { isSnippet, snip, type Snippet } from '../data/snippet.ts'; +import { + isEphemeralSnippet, + isSnippet, + type ResolvedSnippet, + snip, + type Snippet, +} from '../data/snippet.ts'; import { type AnyWgslData, type F32, type I32, isMatInstance, + isNaturallyEphemeral, isVecInstance, WORKAROUND_getSchema, } from '../data/wgslTypes.ts'; @@ -14,8 +22,11 @@ import { type FunctionScopeLayer, getOwnSnippet, type ResolutionCtx, + type SelfResolvable, } from '../types.ts'; import type { ShelllessRepository } from './shellless.ts'; +import { stitch } from '../../src/core/resolve/stitch.ts'; +import { WgslTypeError } from '../../src/errors.ts'; export function numericLiteralToSnippet(value: number): Snippet { if (value >= 2 ** 63 || value < -(2 ** 63)) { @@ -127,3 +138,47 @@ export function coerceToSnippet(value: unknown): Snippet { return snip(value, UnknownData, /* origin */ 'constant'); } + +// defers the resolution of array expressions +export class ArrayExpression implements SelfResolvable { + readonly [$internal] = true; + + constructor( + public readonly elementType: AnyWgslData, + public readonly type: AnyWgslData, + public readonly elements: Snippet[], + ) { + } + + toString(): string { + return 'ArrayExpression'; + } + + [$resolve](ctx: ResolutionCtx): ResolvedSnippet { + for (const elem of this.elements) { + // We check if there are no references among the elements + if ( + (elem.origin === 'argument' && + !isNaturallyEphemeral(elem.dataType)) || + !isEphemeralSnippet(elem) + ) { + const snippetStr = ctx.resolve(elem.value, elem.dataType).value; + const snippetType = + ctx.resolve(concretize(elem.dataType as AnyData)).value; + throw new WgslTypeError( + `'${snippetStr}' reference cannot be used in an array constructor.\n-----\nTry '${snippetType}(${snippetStr})' or 'arrayOf(${snippetType}, count)([...])' to copy the value instead.\n-----`, + ); + } + } + + const arrayType = `array<${ + ctx.resolve(this.elementType).value + }, ${this.elements.length}>`; + + return snip( + stitch`${arrayType}(${this.elements})`, + this.type, + /* origin */ 'runtime', + ); + } +} diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 2c95b96d9e..b47ef79dea 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -33,6 +33,7 @@ import { tryConvertSnippet, } from './conversion.ts'; import { + ArrayExpression, concretize, type GenerationCtx, numericLiteralToSnippet, @@ -498,21 +499,21 @@ ${this.ctx.pre}}`; const [_, calleeNode, argNodes] = expression; const callee = this.expression(calleeNode); - if (wgsl.isWgslStruct(callee.value) || wgsl.isWgslArray(callee.value)) { - // Struct/array schema call. + if (wgsl.isWgslStruct(callee.value)) { + // Struct schema call. if (argNodes.length > 1) { throw new WgslTypeError( - 'Array and struct schemas should always be called with at most 1 argument', + 'Struct schemas should always be called with at most 1 argument', ); } // No arguments `Struct()`, resolve struct name and return. if (!argNodes[0]) { - // the schema becomes the data type + // The schema becomes the data type. return snip( `${this.ctx.resolve(callee.value).value}()`, callee.value, - // A new struct, so not a reference + // A new struct, so not a reference. /* origin */ 'runtime', ); } @@ -527,7 +528,53 @@ ${this.ctx.pre}}`; return snip( this.ctx.resolve(arg.value, callee.value).value, callee.value, - // A new struct, so not a reference + // A new struct, so not a reference. + /* origin */ 'runtime', + ); + } + + if (wgsl.isWgslArray(callee.value)) { + // Array schema call. + if (argNodes.length > 1) { + throw new WgslTypeError( + 'Array schemas should always be called with at most 1 argument', + ); + } + + // No arguments `array<...>()`, resolve array type and return. + if (!argNodes[0]) { + // The schema becomes the data type. + return snip( + `${this.ctx.resolve(callee.value).value}()`, + callee.value, + // A new array, so not a reference. + /* origin */ 'runtime', + ); + } + + const arg = this.typedExpression( + argNodes[0], + callee.value, + ); + + // `d.arrayOf(...)([...])`. + // We resolve each element separately. + if (arg.value instanceof ArrayExpression) { + return snip( + stitch`${ + this.ctx.resolve(callee.value).value + }(${arg.value.elements})`, + arg.dataType, + /* origin */ 'runtime', + ); + } + + // `d.arrayOf(...)(otherArr)`. + // We just let the argument resolve everything. + return snip( + this.ctx.resolve(arg.value, callee.value).value, + callee.value, + // A new array, so not a reference. /* origin */ 'runtime', ); } @@ -720,24 +767,9 @@ ${this.ctx.pre}}`; } } else { // The array is not typed, so we try to guess the types. - const valuesSnippets = valueNodes.map((value) => { - const snippet = this.expression(value as tinyest.Expression); - // We check if there are no references among the elements - if ( - (snippet.origin === 'argument' && - !wgsl.isNaturallyEphemeral(snippet.dataType)) || - !isEphemeralSnippet(snippet) - ) { - const snippetStr = - this.ctx.resolve(snippet.value, snippet.dataType).value; - const snippetType = - this.ctx.resolve(concretize(snippet.dataType as AnyData)).value; - throw new WgslTypeError( - `'${snippetStr}' reference cannot be used in an array constructor.\n-----\nTry '${snippetType}(${snippetStr})' or 'arrayOf(${snippetType}, count)([...])' to copy the value instead.\n-----`, - ); - } - return snippet; - }); + const valuesSnippets = valueNodes.map((value) => + this.expression(value as tinyest.Expression) + ); if (valuesSnippets.length === 0) { throw new WgslTypeError( @@ -756,16 +788,18 @@ ${this.ctx.pre}}`; elemType = concretize(values[0]?.dataType as wgsl.AnyWgslData); } - const arrayType = `array<${ - this.ctx.resolve(elemType).value - }, ${values.length}>`; + const arrayType = arrayOf[$internal].jsImpl( + elemType as wgsl.AnyWgslData, + values.length, + ); return snip( - stitch`${arrayType}(${values})`, - arrayOf[$internal].jsImpl( + new ArrayExpression( elemType as wgsl.AnyWgslData, - values.length, - ) as wgsl.AnyWgslData, + arrayType, + values, + ), + arrayType, /* origin */ 'runtime', ); } diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts index 239ea5ad46..d66c53d739 100644 --- a/packages/typegpu/src/types.ts +++ b/packages/typegpu/src/types.ts @@ -109,15 +109,29 @@ export type FunctionScopeLayer = { reportedReturnTypes: Set; }; +export type SlotBindingLayer = { + type: 'slotBinding'; + bindingMap: WeakMap, unknown>; +}; + +export type BlockScopeLayer = { + type: 'blockScope'; + declarations: Map; +}; + +export type StackLayer = + | ItemLayer + | SlotBindingLayer + | FunctionScopeLayer + | BlockScopeLayer; + export interface ItemStateStack { readonly itemDepth: number; readonly topItem: ItemLayer; readonly topFunctionScope: FunctionScopeLayer | undefined; pushItem(): void; - popItem(): void; pushSlotBindings(pairs: SlotValuePair[]): void; - popSlotBindings(): void; pushFunctionScope( functionType: 'normal' | TgpuShaderStage, args: Snippet[], @@ -129,10 +143,11 @@ export interface ItemStateStack { returnType: AnyData | undefined, externalMap: Record, ): FunctionScopeLayer; - popFunctionScope(): void; pushBlockScope(): void; - popBlockScope(): void; - pop(type?: 'functionScope' | 'blockScope' | 'slotBinding' | 'item'): void; + + pop(type: T): Extract; + pop(): StackLayer | undefined; + readSlot(slot: TgpuSlot): T | undefined; getSnippetById(id: string): Snippet | undefined; defineBlockVariable(id: string, snippet: Snippet): void; diff --git a/packages/typegpu/tests/array.test.ts b/packages/typegpu/tests/array.test.ts index e8220da979..b193c85a20 100644 --- a/packages/typegpu/tests/array.test.ts +++ b/packages/typegpu/tests/array.test.ts @@ -140,6 +140,24 @@ describe('array', () => { ); }); + it('throws when invalid number of arguments during code generation', () => { + const ArraySchema = d.arrayOf(d.u32, 2); + + const f = () => { + 'use gpu'; + // @ts-expect-error + const arr = ArraySchema([1, 1], [6, 7]); + return; + }; + + expect(() => tgpu.resolve([f])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn*:f + - fn*:f(): Array schemas should always be called with at most 1 argument] + `); + }); + it('can be called to create a default value', () => { const ArraySchema = d.arrayOf(d.vec3f, 2); @@ -188,16 +206,28 @@ describe('array', () => { it('generates correct code when array clone is used', () => { const ArraySchema = d.arrayOf(d.u32, 1); - const testFn = tgpu.fn([])(() => { + const f = (arr: d.Infer) => { + 'use gpu'; + const clone = ArraySchema(arr); + }; + + const testFn = () => { + 'use gpu'; const myArray = ArraySchema([d.u32(10)]); const myClone = ArraySchema(myArray); + f(myArray); return; - }); + }; expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` - "fn testFn() { + "fn f(arr: array) { + var clone = arr; + } + + fn testFn() { var myArray = array(10u); var myClone = myArray; + f(myArray); return; }" `); @@ -221,6 +251,65 @@ describe('array', () => { `); }); + it('generates correct code when array expression with ephemeral element type clone is used', () => { + const f = () => { + 'use gpu'; + const arr = d.arrayOf(d.f32, 2)([6, 7]); + return; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() { + var arr = array(6f, 7f); + return; + }" + `); + }); + + it('generates correct code when array expression with reference element type clone is used', () => { + const f = (v: d.v4f) => { + 'use gpu'; + const v2 = d.vec4f(3); + const v3 = v2; + const arr = d.arrayOf(d.vec4f, 3)([v, v2, v3]); + }; + + const main = tgpu.fn([])(() => { + const v1 = d.vec4f(7); + f(v1); + return; + }); + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "fn f(v: vec4f) { + var v2 = vec4f(3); + let v3 = (&v2); + var arr = array(v, v2, (*v3)); + } + + fn main() { + var v1 = vec4f(7); + f(v1); + return; + }" + `); + }); + + it('generates correct code when array expression with mixed element types clone is used', () => { + const f = () => { + 'use gpu'; + const arr = d.arrayOf(d.f32, 3)([5, 6.7, 8.0]); + return; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() { + var arr = array(5f, 6.7f, 8f); + return; + }" + `); + }); + it('can be immediately-invoked in TGSL', () => { const foo = tgpu.fn([])(() => { const result = d.arrayOf(d.f32, 4)(); @@ -339,7 +428,8 @@ describe('array', () => { expect(() => tgpu.resolve([foo])).toThrowErrorMatchingInlineSnapshot(` [Error: Resolution of the following tree failed: - - - fn:foo: 'myVec' reference cannot be used in an array constructor. + - fn:foo + - ArrayExpression: 'myVec' reference cannot be used in an array constructor. ----- Try 'vec2f(myVec)' or 'arrayOf(vec2f, count)([...])' to copy the value instead. -----] @@ -354,7 +444,8 @@ describe('array', () => { expect(() => tgpu.resolve([foo])).toThrowErrorMatchingInlineSnapshot(` [Error: Resolution of the following tree failed: - - - fn:foo: 'myVec' reference cannot be used in an array constructor. + - fn:foo + - ArrayExpression: 'myVec' reference cannot be used in an array constructor. ----- Try 'vec2f(myVec)' or 'arrayOf(vec2f, count)([...])' to copy the value instead. -----] diff --git a/packages/typegpu/tests/slot.test.ts b/packages/typegpu/tests/slot.test.ts index 54ee250267..0d222143e0 100644 --- a/packages/typegpu/tests/slot.test.ts +++ b/packages/typegpu/tests/slot.test.ts @@ -1,7 +1,7 @@ import { describe, expect } from 'vitest'; import * as d from '../src/data/index.ts'; -import * as std from '../src/std/index.ts'; import tgpu from '../src/index.ts'; +import * as std from '../src/std/index.ts'; import { it } from './utils/extendedIt.ts'; const RED = 'vec3f(1., 0., 0.)'; @@ -363,4 +363,32 @@ describe('tgpu.slot', () => { }" `); }); + + it('includes slot bindings in toString', () => { + const firstSlot = tgpu.slot(); + const secondSlot = tgpu.slot(); + const thirdSlot = tgpu.slot(); + + const getSize = tgpu.fn([], d.f32)(() => + firstSlot.$ + secondSlot.$ + thirdSlot.$ + ) + .with(firstSlot, 1) + .with(secondSlot, 2) + .with(thirdSlot, 3); + + expect(getSize.toString()).toMatchInlineSnapshot( + `"fn:getSize[firstSlot=1, secondSlot=2, thirdSlot=3]"`, + ); + }); + + it('safe stringifies in toString', () => { + const slot = tgpu.slot(); + + const getSize = tgpu.fn([], d.f32)(() => slot.$.x) + .with(slot, d.vec4f(1, 2, 3, 4)); + + expect(getSize.toString()).toMatchInlineSnapshot( + `"fn:getSize[slot=vec4f(1, 2, 3, 4)]"`, + ); + }); }); diff --git a/packages/typegpu/tests/struct.test.ts b/packages/typegpu/tests/struct.test.ts index b7fe34dfb2..54d644c4a5 100644 --- a/packages/typegpu/tests/struct.test.ts +++ b/packages/typegpu/tests/struct.test.ts @@ -390,6 +390,29 @@ describe('struct', () => { ); }); + it('throws when invalid number of arguments during code generation', () => { + const Boid = struct({ + pos: vec2f, + vel: vec2f, + }); + + const f = () => { + 'use gpu'; + const b1 = Boid({ pos: vec2f(6), vel: vec2f(7) }); + + // @ts-expect-error + const b2 = Boid(b1, b1); + return; + }; + + expect(() => tgpu.resolve([f])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn*:f + - fn*:f(): Struct schemas should always be called with at most 1 argument] + `); + }); + it('allows builtin names as struct props', () => { const myStruct = struct({ min: u32, diff --git a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts index 3c35557242..bee025f1f1 100644 --- a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts +++ b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts @@ -14,6 +14,7 @@ import * as std from '../../src/std/index.ts'; import wgslGenerator from '../../src/tgsl/wgslGenerator.ts'; import { CodegenState } from '../../src/types.ts'; import { it } from '../utils/extendedIt.ts'; +import { ArrayExpression } from '../../src/tgsl/generationHelpers.ts'; const { NodeTypeCatalog: NODE } = tinyest; @@ -338,7 +339,7 @@ describe('wgslGenerator', () => { const res2 = wgslGenerator.expression( (astInfo.ast?.body[1][1] as tinyest.Const)[2], ); - ctx[$internal].itemStateStack.popBlockScope(); + ctx[$internal].itemStateStack.pop('blockScope'); expect(res2.dataType).toStrictEqual(d.vec4f); @@ -353,7 +354,7 @@ describe('wgslGenerator', () => { const res4 = wgslGenerator.expression( astInfo.ast?.body[1][2] as tinyest.Expression, ); - ctx[$internal].itemStateStack.popBlockScope(); + ctx[$internal].itemStateStack.pop('blockScope'); expect(res3.dataType).toStrictEqual(d.atomic(d.u32)); expect(res4.dataType).toStrictEqual(Void); @@ -528,6 +529,44 @@ describe('wgslGenerator', () => { }); }); + it('creates intermediate representation for array expression', () => { + const testFn = tgpu.fn([], d.u32)(() => { + const arr = [d.u32(1), 2, 3]; + return arr[1] as number; + }); + + const astInfo = getMetaData( + testFn[$internal].implementation as (...args: unknown[]) => unknown, + ); + + if (!astInfo) { + throw new Error('Expected prebuilt AST to be present'); + } + + provideCtx(ctx, () => { + ctx[$internal].itemStateStack.pushFunctionScope( + 'normal', + [], + {}, + d.u32, + (astInfo.externals as () => Record)() ?? {}, + ); + + wgslGenerator.initGenerator(ctx); + const res = wgslGenerator.expression( + // deno-fmt-ignore: it's better that way + ( + astInfo.ast?.body[1][0] as tinyest.Const + )[2] as unknown as tinyest.Expression, + ); + + expect(res.value instanceof ArrayExpression).toBe(true); + expect((res.value as unknown as ArrayExpression).type).toBe(res.dataType); + expect((res.value as unknown as ArrayExpression).elementType) + .toBe((res.dataType as unknown as WgslArray).elementType); + }); + }); + it('generates correct code for array expressions', () => { const testFn = tgpu.fn([], d.u32)(() => { const arr = [d.u32(1), 2, 3]; @@ -594,6 +633,42 @@ describe('wgslGenerator', () => { return arr[1i].x; }" `); + + const astInfo = getMetaData( + testFn[$internal].implementation as (...args: unknown[]) => unknown, + ); + + if (!astInfo) { + throw new Error('Expected prebuilt AST to be present'); + } + + expect(JSON.stringify(astInfo.ast?.body)).toMatchInlineSnapshot( + `"[0,[[13,"arr",[100,[[6,[7,"d","vec2u"],[[5,"1"],[5,"2"]]],[6,[7,"d","vec2u"],[[5,"3"],[5,"4"]]],[6,[7,"std","min"],[[6,[7,"d","vec2u"],[[5,"5"],[5,"8"]]],[6,[7,"d","vec2u"],[[5,"7"],[5,"6"]]]]]]]],[10,[7,[8,"arr",[5,"1"]],"x"]]]]"`, + ); + + provideCtx(ctx, () => { + ctx[$internal].itemStateStack.pushFunctionScope( + 'normal', + [], + {}, + d.u32, + (astInfo.externals as () => Record)() ?? {}, + ); + + // Check for: const arr = [1, 2, 3] + // ^ this should be an array + wgslGenerator.initGenerator(ctx); + const res = wgslGenerator.expression( + // deno-fmt-ignore: it's better that way + ( + astInfo.ast?.body[1][0] as tinyest.Const + )[2] as unknown as tinyest.Expression, + ); + + expect(d.isWgslArray(res.dataType)).toBe(true); + expect((res.dataType as unknown as WgslArray).elementCount).toBe(3); + expect((res.dataType as unknown as WgslArray).elementType).toBe(d.vec2u); + }); }); it('does not autocast lhs of an assignment', () => { @@ -696,6 +771,25 @@ describe('wgslGenerator', () => { expect(JSON.stringify(astInfo.ast?.body)).toMatchInlineSnapshot( `"[0,[[13,"arr",[100,[[7,"derivedV2f","$"],[6,[7,"std","mul"],[[7,"derivedV2f","$"],[6,[7,"d","vec2f"],[[5,"2"],[5,"2"]]]]]]]],[10,[7,[8,"arr",[5,"1"]],"y"]]]]"`, ); + + const res = provideCtx(ctx, () => { + ctx[$internal].itemStateStack.pushFunctionScope( + 'normal', + [], + {}, + d.f32, + (astInfo.externals as () => Record)() ?? {}, + ); + + wgslGenerator.initGenerator(ctx); + return wgslGenerator.expression( + (astInfo.ast?.body[1][0] as tinyest.Const)[2] as tinyest.Expression, + ); + }); + + expect(d.isWgslArray(res.dataType)).toBe(true); + expect((res.dataType as unknown as WgslArray).elementCount).toBe(2); + expect((res.dataType as unknown as WgslArray).elementType).toBe(d.vec2f); }); it('allows for member access on values returned from function calls', () => { diff --git a/packages/typegpu/tests/utils/parseResolved.ts b/packages/typegpu/tests/utils/parseResolved.ts index fc59207de0..515faf4fe0 100644 --- a/packages/typegpu/tests/utils/parseResolved.ts +++ b/packages/typegpu/tests/utils/parseResolved.ts @@ -58,8 +58,8 @@ function extractSnippetFromFn(cb: () => unknown): Snippet { } finally { if (pushedFnScope) { ctx.popBlockScope(); - ctx[$internal].itemStateStack.popFunctionScope(); - ctx[$internal].itemStateStack.popItem(); + ctx[$internal].itemStateStack.pop('functionScope'); + ctx[$internal].itemStateStack.pop('item'); } ctx.popMode('codegen'); }