diff --git a/packages/typegpu/src/tgsl/arrayExpression.ts b/packages/typegpu/src/tgsl/arrayExpression.ts new file mode 100644 index 0000000000..54665614bb --- /dev/null +++ b/packages/typegpu/src/tgsl/arrayExpression.ts @@ -0,0 +1,61 @@ +import { isStoredInMemory } from '../data/snippet.ts'; +import type { ResolvedSnippet, Snippet } from '../data/snippet.ts'; +import { + type AnyWgslData, + type BaseData, + isNaturallyEphemeral, + type WgslArray, +} from '../data/wgslTypes.ts'; +import { WgslTypeError } from '../errors.ts'; +import { $comptimeValueOf, $internal, $resolve } from '../shared/symbols.ts'; +import { + asComptime, + type ComptimeValue, + type ResolutionCtx, + type SelfResolvable, + type WithComptimeValue, +} from '../types.ts'; +import { stringifySnippet } from './stringifySnippet.ts'; + +export class ArrayExpression implements SelfResolvable, WithComptimeValue { + readonly [$internal] = true; + readonly type: WgslArray; + readonly elements: Snippet[]; + + constructor(type: WgslArray, elements: Snippet[]) { + this.type = type; + this.elements = elements; + } + + toString(): string { + return 'ArrayExpression'; + } + + [$comptimeValueOf](): ComptimeValue { + // The array is known at comptime only if every element is known at comptime + const values: unknown[] = []; + for (const elem of this.elements) { + const known = asComptime(elem); + if (!known) { + return { comptime: false }; + } + values.push(known.value); + } + return { comptime: true, value: values }; + } + + [$resolve](ctx: ResolutionCtx): ResolvedSnippet { + for (const elem of this.elements) { + // We check if there are no references among the elements + if (isStoredInMemory(elem) && !isNaturallyEphemeral(elem.dataType)) { + const snippetStr = stringifySnippet(elem); + const snippetType = ctx.resolve(elem.dataType as BaseData).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 ctx.gen.typeInstantiation(this.type, this.elements); + } +} diff --git a/packages/typegpu/src/tgsl/conversion.ts b/packages/typegpu/src/tgsl/conversion.ts index 1a84ac3d7d..21bdc9fd06 100644 --- a/packages/typegpu/src/tgsl/conversion.ts +++ b/packages/typegpu/src/tgsl/conversion.ts @@ -13,9 +13,11 @@ import { isMat, isPtr, isVec, + isWgslArray, isWgslStruct, type Ptr, type U32, + type WgslArray, type WgslStruct, } from '../data/wgslTypes.ts'; import { invariant, WgslTypeError } from '../errors.ts'; @@ -25,6 +27,7 @@ import { assertExhaustive } from '../shared/utilityTypes.ts'; import { logger } from '../tgpuLogger.ts'; import type { ResolutionCtx } from '../types.ts'; import { accessStructProp } from './accessStructProp.ts'; +import { ArrayExpression } from './arrayExpression.ts'; import { stringifySnippet } from './stringifySnippet.ts'; type ConversionAction = 'ref' | 'deref' | 'cast' | 'none'; @@ -38,6 +41,19 @@ const INFINITE_RANK: ConversionRankInfo = { action: 'none', }; +function needsArrayCast(src: WgslArray, dest: WgslArray): boolean { + // An element count of 0 means runtime-sized, which can't be constructed element by element. + if (src === dest || src.elementCount === 0 || dest.elementCount === 0) { + return false; + } + const elementRank = getAutoConversionRank(src.elementType, dest.elementType); + return ( + src.elementCount !== dest.elementCount || + elementRank.rank === Number.POSITIVE_INFINITY || + elementRank.action !== 'none' + ); +} + function getAutoConversionRank(src: BaseData, dest: BaseData): ConversionRankInfo { const trueSrc = undecorate(src); const trueDst = undecorate(dest); @@ -46,6 +62,12 @@ function getAutoConversionRank(src: BaseData, dest: BaseData): ConversionRankInf if (trueSrc.type === 'struct' && trueSrc !== trueDst) { return { rank: 1, action: 'cast', targetType: dest }; } + if (isWgslArray(trueSrc) && isWgslArray(trueDst) && needsArrayCast(trueSrc, trueDst)) { + // Includes the cost of converting an element, so that `[[1], [1.5]]` becomes f32 in any row order. + const elementRank = getConversionRank(trueSrc.elementType, trueDst.elementType, true).rank; + const rank = 1 + (Number.isFinite(elementRank) ? elementRank : 0); + return { rank, action: 'cast', targetType: dest }; + } return { rank: 0, action: 'none' }; } @@ -263,8 +285,12 @@ function applyActionToSnippet( snippet.origin, snippet.possibleSideEffects, ); - case 'deref': - return derefSnippet(snippet); + case 'deref': { + const dereferenced = derefSnippet(snippet); + return isWgslArray(targetType) && isWgslArray(dereferenced.dataType) + ? tryConvertSnippet(ctx, dereferenced, targetType) + : dereferenced; + } case 'cast': { if (isWgslStruct(snippet.dataType) && isWgslStruct(targetType)) { const typeName = getName(snippet.dataType) ?? ''; @@ -300,6 +326,59 @@ function applyActionToSnippet( false, ); } + if (isWgslArray(snippet.dataType) && isWgslArray(targetType)) { + const typeName = String(snippet.dataType); + const targetName = String(targetType); + + const convertElement = (element: Snippet) => { + const converted = convertToCommonType(ctx, [element], [targetType.elementType]); + if (!converted || !converted[0]) { + throw new Error( + `Cannot auto-convert array '${typeName}' to '${targetName}' because its elements are not convertible.`, + ); + } + return converted[0]; + }; + + // Array to array casting + if (snippet.dataType.elementCount !== targetType.elementCount) { + throw new Error( + `Cannot auto-convert array '${typeName}' to '${targetName}' because their element counts differ.`, + ); + } + + // An array literal stays one, so that its elements can still be known at comptime. + if (snippet.value instanceof ArrayExpression) { + return snip( + new ArrayExpression( + targetType as WgslArray, + snippet.value.elements.map(convertElement), + ), + targetType, + snippet.origin, + snippet.possibleSideEffects, + ); + } + + if (snippet.possibleSideEffects) { + throw new Error( + `Cannot resolve array cast from '${typeName}' to '${targetName}'. Store the value to a variable first, then cast it.`, + ); + } + + const source = ctx.resolveSnippet(snippet).value; + const elementType = snippet.dataType.elementType; + const elementSnips = Array.from({ length: targetType.elementCount }, (_, index) => + convertElement(snip(`${source}[${index}]`, elementType, 'runtime', false)), + ); + const targetSnippet = ctx.resolve(targetType).value; + return snip( + `${targetSnippet}(${elementSnips.map((elem) => ctx.resolveSnippet(elem).value).join(', ')})`, + targetType, + 'runtime', + false, + ); + } // Casting means calling the schema with the snippet as an argument. return schemaCallWrapperGPU(ctx, targetType, snippet); } diff --git a/packages/typegpu/src/tgsl/generationHelpers.ts b/packages/typegpu/src/tgsl/generationHelpers.ts index d62bc4a361..569c2747fd 100644 --- a/packages/typegpu/src/tgsl/generationHelpers.ts +++ b/packages/typegpu/src/tgsl/generationHelpers.ts @@ -1,31 +1,21 @@ import { UnknownData } from '../data/dataTypes.ts'; import { abstractFloat, abstractInt, bool, f32, i32 } from '../data/numeric.ts'; import { isRef } from '../data/ref.ts'; -import { isSnippet, isStoredInMemory, snip, withDataType } from '../data/snippet.ts'; -import type { ResolvedSnippet, Snippet } from '../data/snippet.ts'; +import { isSnippet, snip, withDataType } from '../data/snippet.ts'; +import type { Snippet } from '../data/snippet.ts'; import { type AnyWgslData, type BaseData, type F32, type I32, isMatInstance, - isNaturallyEphemeral, isVecInstance, - type WgslArray, WORKAROUND_getSchema, } from '../data/wgslTypes.ts'; -import { - asComptime, - type ComptimeValue, - getOwnSnippet, - type ResolutionCtx, - type SelfResolvable, - type WithComptimeValue, -} from '../types.ts'; -import { WgslTypeError } from '../errors.ts'; -import { $comptimeValueOf, $internal, $resolve } from '../shared/symbols.ts'; +import { getOwnSnippet } from '../types.ts'; import { logger } from '../tgpuLogger.ts'; -import { stringifySnippet } from './stringifySnippet.ts'; + +export { ArrayExpression } from './arrayExpression.ts'; export function numericLiteralToSnippet(value: number): Snippet { if (value >= 2 ** 63 || value < -(2 ** 63)) { @@ -106,45 +96,3 @@ export function coerceToSnippet(value: unknown): Snippet { * Defers resolution. Stores array elements as snippets so the * generator can access them when needed. */ -export class ArrayExpression implements SelfResolvable, WithComptimeValue { - readonly [$internal] = true; - readonly type: WgslArray; - readonly elements: Snippet[]; - - constructor(type: WgslArray, elements: Snippet[]) { - this.type = type; - this.elements = elements; - } - - toString(): string { - return 'ArrayExpression'; - } - - [$comptimeValueOf](): ComptimeValue { - // The array is known at comptime only if every element is known at comptime - const values: unknown[] = []; - for (const elem of this.elements) { - const known = asComptime(elem); - if (!known) { - return { comptime: false }; - } - values.push(known.value); - } - return { comptime: true, value: values }; - } - - [$resolve](ctx: ResolutionCtx): ResolvedSnippet { - for (const elem of this.elements) { - // We check if there are no references among the elements - if (isStoredInMemory(elem) && !isNaturallyEphemeral(elem.dataType)) { - const snippetStr = stringifySnippet(elem); - const snippetType = ctx.resolve(concretize(elem.dataType as BaseData)).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 ctx.gen.typeInstantiation(this.type, this.elements); - } -} diff --git a/packages/typegpu/tests/array.test.ts b/packages/typegpu/tests/array.test.ts index f287f22f7e..74e2887a71 100644 --- a/packages/typegpu/tests/array.test.ts +++ b/packages/typegpu/tests/array.test.ts @@ -379,6 +379,128 @@ describe('array', () => { `); }); + it('converts the elements when cast to an array of another element type', () => { + const main = () => { + 'use gpu'; + const arr = [1.5, 2.5, 3.5]; + const b = d.arrayOf(d.u32, 3)(arr); + }; + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "fn main() { + let arr = array(1.5, 2.5, 3.5); + let b = array(u32(arr[0]), u32(arr[1]), u32(arr[2])); + }" + `); + }); + + it('converts the elements when passed as an argument of another element type', () => { + const A = d.arrayOf(d.struct({ p: d.u32 }), 2); + const B = d.arrayOf(d.struct({ p: d.f32 }), 2); + const helper = tgpu.fn([B])(() => {}); + + const main = () => { + 'use gpu'; + const a = A(); + helper(a); + }; + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "struct item { + p: u32, + } + + struct item_1 { + p: f32, + } + + fn helper(_arg_0: array) {} + + fn main() { + let a = array(); + helper(array(item_1(f32(a[0].p)), item_1(f32(a[1].p)))); + }" + `); + }); + + it('converts the elements of an alias', () => { + const main = () => { + 'use gpu'; + const arr = [1.5, 2.5]; + const alias = arr; + const b = d.arrayOf(d.u32, 2)(alias); + }; + + expect(tgpu.resolve([main])).toContain( + 'let b = array(u32((*alias_1)[0]), u32((*alias_1)[1]));', + ); + }); + + it('throws when an alias is cast to an array of another length', () => { + const helper = tgpu.fn([d.arrayOf(d.f32, 5)])(() => {}); + + const main = () => { + 'use gpu'; + const arr = [1.5, 2.5]; + const alias = arr; + helper(alias); + }; + + expect(() => tgpu.resolve([main])).toThrow(/because their element counts differ/); + }); + + it('infers the same nested array type in any row order', () => { + const first = () => { + 'use gpu'; + const a = [[1], [1.5]]; + }; + const second = () => { + 'use gpu'; + const a = [[1.5], [1]]; + }; + + expect(tgpu.resolve([first])).toContain('array, 2>'); + expect(tgpu.resolve([second])).toContain('array, 2>'); + expect(tgpu.resolve([first])).toContain('1.5'); + }); + + it('does not cast between separate schemas of the same array type', () => { + const A = d.arrayOf(d.i32, 2); + const B = d.arrayOf(d.i32, 2); + const helper = tgpu.fn([B])(() => {}); + + const main = () => { + 'use gpu'; + const a = A(); + helper(a); + }; + + expect(tgpu.resolve([main])).toContain('helper(a);'); + }); + + it('throws when cast to an array of another length', () => { + const helper = tgpu.fn([d.arrayOf(d.i32, 5)])(() => {}); + + const main = () => { + 'use gpu'; + const arr = [1, 2, 3]; + return helper(arr); + }; + + expect(() => tgpu.resolve([main])).toThrow(/because their element counts differ/); + }); + + it('throws when cast from an expression with possible side effects', () => { + const get = tgpu.fn([], d.arrayOf(d.f32, 2))(() => d.arrayOf(d.f32, 2)()); + + const main = () => { + 'use gpu'; + const b = d.arrayOf(d.u32, 2)(get()); + }; + + expect(() => tgpu.resolve([main])).toThrow(/Store the value to a variable first/); + }); + it('can be immediately-partially-invoked and initialized in TGSL', () => { const foo = tgpu.fn([])(() => { const result = d.arrayOf(d.f32)(4)([4, 3, 2, 1]);