diff --git a/packages/typegpu/src/core/function/tgpuFn.ts b/packages/typegpu/src/core/function/tgpuFn.ts index 8a5c9bd3ae..d8cfa71421 100644 --- a/packages/typegpu/src/core/function/tgpuFn.ts +++ b/packages/typegpu/src/core/function/tgpuFn.ts @@ -145,8 +145,8 @@ export function isTgpuFn( value: unknown, ): value is TgpuFn<(...args: Args) => Return> { return ( - isMarkedInternal(value) && - (value as TgpuFn<(...args: Args) => Return>)?.resourceType === 'function' + (value as TgpuFn<(...args: Args) => Return>)?.resourceType === 'function' && + isMarkedInternal(value) ); } @@ -154,8 +154,8 @@ export function isGenericFn( value: unknown, ): value is TgpuGenericFn { return ( - isMarkedInternal(value) && - (value as TgpuGenericFn)?.resourceType === 'generic-function' + (value as TgpuGenericFn)?.resourceType === 'generic-function' && + isMarkedInternal(value) ); } diff --git a/packages/typegpu/src/core/resolve/externals.ts b/packages/typegpu/src/core/resolve/externals.ts index 4d25e11068..05fca8f375 100644 --- a/packages/typegpu/src/core/resolve/externals.ts +++ b/packages/typegpu/src/core/resolve/externals.ts @@ -20,6 +20,10 @@ function isResolvable(value: unknown) { */ export function mergeFunctionExternals(fnExternals: FnExternals): ExternalMap { const base = fnExternals.pluginProvided ?? fnExternals.userProvided ?? {}; + if (fnExternals.args === undefined && fnExternals.out === undefined) { + // Nothing to merge in, and consumers never mutate the map + return base; + } // avoid calling any of the getters const result: ExternalMap = Object.defineProperties({}, Object.getOwnPropertyDescriptors(base)); for (const flatExternal of [fnExternals.args, fnExternals.out].filter((e) => e !== undefined)) { diff --git a/packages/typegpu/src/core/resolve/namespace.ts b/packages/typegpu/src/core/resolve/namespace.ts index a20f15b6ab..fb50f0a212 100644 --- a/packages/typegpu/src/core/resolve/namespace.ts +++ b/packages/typegpu/src/core/resolve/namespace.ts @@ -1,5 +1,4 @@ import type { ResolvedSnippet } from '../../data/snippet.ts'; -import { bannedTokens, builtins, templateEnumerants } from '../../nameUtils.ts'; import { $internal } from '../../shared/symbols.ts'; import { ShelllessRepository } from '../../tgsl/shellless.ts'; import type { TgpuLazy, TgpuSlot } from '../slot/slotTypes.ts'; @@ -7,6 +6,10 @@ import type { TgpuLazy, TgpuSlot } from '../slot/slotTypes.ts'; type SlotToValueMap = Map, unknown>; export interface NamespaceInternal { + /** + * Global identifiers taken by resolved items. Reserved WGSL words are checked separately, + * so that creating a namespace doesn't have to copy them. + */ readonly takenGlobalIdentifiers: Set; readonly shelllessRepo: ShelllessRepository; readonly strategy: 'random' | 'strict'; @@ -36,7 +39,7 @@ class NamespaceImpl implements Namespace { constructor(strategy: 'random' | 'strict') { this[$internal] = { strategy, - takenGlobalIdentifiers: new Set([...bannedTokens, ...builtins, ...templateEnumerants]), + takenGlobalIdentifiers: new Set(), shelllessRepo: new ShelllessRepository(), memoizedResolves: new WeakMap(), memoizedLazy: new WeakMap(), diff --git a/packages/typegpu/src/core/resolve/stitch.ts b/packages/typegpu/src/core/resolve/stitch.ts index b9664d8c3b..a5886054e3 100644 --- a/packages/typegpu/src/core/resolve/stitch.ts +++ b/packages/typegpu/src/core/resolve/stitch.ts @@ -4,6 +4,10 @@ import type { ResolutionCtx } from '../../types.ts'; type ValueOrArray = T | readonly T[]; +function resolvePart(ctx: ResolutionCtx, part: Snippet | string | number) { + return isSnippet(part) ? ctx.resolveSnippet(part).value : part; +} + /** * "The reverse of snipping" * Injects resolved snippets into a template string. @@ -14,21 +18,22 @@ export function stitch( ) { const ctx = getResolutionCtx() as ResolutionCtx; - function resolveSnippet(maybeSnippet: Snippet | string | number) { - return isSnippet(maybeSnippet) ? ctx.resolveSnippet(maybeSnippet).value : maybeSnippet; - } - let result = ''; for (let i = 0; i < strings.length; ++i) { result += strings[i]; const snippet = snippets[i]; if (Array.isArray(snippet)) { - result += snippet - .filter((s) => s !== undefined) - .map(resolveSnippet) - .join(', '); + // Joining the defined elements with ', ' + let first = true; + for (const element of snippet) { + if (element === undefined) { + continue; + } + result += (first ? '' : ', ') + resolvePart(ctx, element); + first = false; + } } else if (snippet) { - result += resolveSnippet(snippet as Snippet | string | number); + result += resolvePart(ctx, snippet as Snippet | string | number); } } return result; diff --git a/packages/typegpu/src/core/texture/texture.ts b/packages/typegpu/src/core/texture/texture.ts index 48e46a9f71..f6d6b83e35 100644 --- a/packages/typegpu/src/core/texture/texture.ts +++ b/packages/typegpu/src/core/texture/texture.ts @@ -743,10 +743,15 @@ class TgpuFixedTextureViewImpl return this; } + // Created once, since shader code can access a view many times + #valueProxy: Infer | undefined; + get [$gpuValueOf](): Infer { + if (this.#valueProxy) { + return this.#valueProxy; + } const schema = this.schema; - - return new Proxy( + this.#valueProxy = new Proxy( { [$internal]: true, get [$ownSnippet]() { @@ -758,6 +763,7 @@ class TgpuFixedTextureViewImpl }, valueProxyHandler, ) as unknown as Infer; + return this.#valueProxy; } get $(): Infer { @@ -837,9 +843,15 @@ export class TgpuLaidOutTextureViewImpl | undefined; + get [$gpuValueOf](): Infer { + if (this.#valueProxy) { + return this.#valueProxy; + } const schema = this.schema; - return new Proxy( + this.#valueProxy = new Proxy( { [$internal]: true, get [$ownSnippet]() { @@ -851,6 +863,7 @@ export class TgpuLaidOutTextureViewImpl; + return this.#valueProxy; } get $(): Infer { diff --git a/packages/typegpu/src/data/dataTypes.ts b/packages/typegpu/src/data/dataTypes.ts index b2b117820b..6e4f1811ff 100644 --- a/packages/typegpu/src/data/dataTypes.ts +++ b/packages/typegpu/src/data/dataTypes.ts @@ -176,7 +176,7 @@ export function isLooseData(data: unknown): data is AnyLooseData { * isDisarray(d.vec3f) // false */ export function isDisarray(schema: unknown): schema is Disarray { - return isMarkedInternal(schema) && (schema as Disarray)?.type === 'disarray'; + return (schema as Disarray)?.type === 'disarray' && isMarkedInternal(schema); } /** @@ -193,11 +193,11 @@ export function isDisarray(schema: unknown): schema is Disarray { * isUnstruct(d.vec3f) // false */ export function isUnstruct(schema: unknown): schema is Unstruct { - return isMarkedInternal(schema) && (schema as Unstruct)?.type === 'unstruct'; + return (schema as Unstruct)?.type === 'unstruct' && isMarkedInternal(schema); } export function isLooseDecorated(value: unknown): value is LooseDecorated { - return isMarkedInternal(value) && (value as LooseDecorated)?.type === 'loose-decorated'; + return (value as LooseDecorated)?.type === 'loose-decorated' && isMarkedInternal(value); } export function getCustomAlignment(data: wgsl.BaseData): number | undefined { diff --git a/packages/typegpu/src/data/vertexFormatData.ts b/packages/typegpu/src/data/vertexFormatData.ts index 5a4a8c517d..b407f22e29 100644 --- a/packages/typegpu/src/data/vertexFormatData.ts +++ b/packages/typegpu/src/data/vertexFormatData.ts @@ -278,7 +278,7 @@ export type PackedData = | unorm8x4_bgra; export function isPackedData(value: unknown): value is PackedData { - return isMarkedInternal(value) && packedFormats.has((value as PackedData)?.type); + return packedFormats.has((value as PackedData)?.type) && isMarkedInternal(value); } type U32Data = U32 | Vec2u | Vec3u | Vec4u; diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts index f39bf80fb6..ffa548c6d4 100644 --- a/packages/typegpu/src/data/wgslTypes.ts +++ b/packages/typegpu/src/data/wgslTypes.ts @@ -1554,7 +1554,7 @@ export type AnyWgslData = export function isVecInstance(value: unknown): value is AnyVecInstance { const v = value as AnyVecInstance | undefined; - return isMarkedInternal(v) && typeof v.kind === 'string' && v.kind.startsWith('vec'); + return typeof v?.kind === 'string' && v.kind.startsWith('vec') && isMarkedInternal(v); } export function isVecBoolInstance(value: unknown): value is v2b | v3b | v4b { @@ -1563,17 +1563,17 @@ export function isVecBoolInstance(value: unknown): value is v2b | v3b | v4b { export function isVec2(value: unknown): value is Vec2f | Vec2h | Vec2i | Vec2u { const v = value as AnyWgslData | undefined; - return isMarkedInternal(v) && typeof v.type === 'string' && v.type.startsWith('vec2'); + return typeof v?.type === 'string' && v.type.startsWith('vec2') && isMarkedInternal(v); } export function isVec3(value: unknown): value is Vec3f | Vec3h | Vec3i | Vec3u { const v = value as AnyWgslData | undefined; - return isMarkedInternal(v) && typeof v.type === 'string' && v.type.startsWith('vec3'); + return typeof v?.type === 'string' && v.type.startsWith('vec3') && isMarkedInternal(v); } export function isVec4(value: unknown): value is Vec4f | Vec4h | Vec4i | Vec4u { const v = value as AnyWgslData | undefined; - return isMarkedInternal(v) && typeof v.type === 'string' && v.type.startsWith('vec4'); + return typeof v?.type === 'string' && v.type.startsWith('vec4') && isMarkedInternal(v); } export function isVec( @@ -1604,20 +1604,20 @@ export function isVecBool(value: unknown): value is Vec2b | Vec3b | Vec4b { export function isMatInstance(value: unknown): value is AnyMatInstance { const v = value as AnyMatInstance | undefined; return ( - isMarkedInternal(v) && typeof v.kind?.startsWith === 'function' && v.kind.startsWith('mat') + typeof v?.kind?.startsWith === 'function' && v.kind.startsWith('mat') && isMarkedInternal(v) ); } export function isMat2x2f(value: unknown): value is Mat2x2f { - return isMarkedInternal(value) && (value as AnyWgslData)?.type === 'mat2x2f'; + return (value as AnyWgslData)?.type === 'mat2x2f' && isMarkedInternal(value); } export function isMat3x3f(value: unknown): value is Mat3x3f { - return isMarkedInternal(value) && (value as AnyWgslData)?.type === 'mat3x3f'; + return (value as AnyWgslData)?.type === 'mat3x3f' && isMarkedInternal(value); } export function isMat4x4f(value: unknown): value is Mat4x4f { - return isMarkedInternal(value) && (value as AnyWgslData)?.type === 'mat4x4f'; + return (value as AnyWgslData)?.type === 'mat4x4f' && isMarkedInternal(value); } export function isMat(value: unknown): value is Mat2x2f | Mat3x3f | Mat4x4f { @@ -1626,7 +1626,7 @@ export function isMat(value: unknown): value is Mat2x2f | Mat3x3f | Mat4x4f { export function isInteger(value: unknown): value is AbstractInt | I32 | U32 { return ( - isMarkedInternal(value) && ['abstractInt', 'i32', 'u32'].includes((value as AnyWgslData)?.type) + ['abstractInt', 'i32', 'u32'].includes((value as AnyWgslData)?.type) && isMarkedInternal(value) ); } @@ -1665,7 +1665,7 @@ export function isWgslData(value: unknown): value is AnyWgslData { * isWgslArray(d.vec3f) // false */ export function isWgslArray(schema: unknown): schema is WgslArray { - return isMarkedInternal(schema) && (schema as WgslArray)?.type === 'array'; + return (schema as WgslArray)?.type === 'array' && isMarkedInternal(schema); } /** @@ -1681,7 +1681,7 @@ export function isWgslArray(schema: unknown): schema is WgslArray { * isWgslStruct(d.vec3f) // false */ export function isWgslStruct(schema: unknown): schema is WgslStruct { - return isMarkedInternal(schema) && (schema as WgslStruct)?.type === 'struct'; + return (schema as WgslStruct)?.type === 'struct' && isMarkedInternal(schema); } /** @@ -1693,7 +1693,7 @@ export function isWgslStruct(schema: unknown): schema is WgslStruct { * isPtr(d.f32) // false */ export function isPtr(schema: unknown): schema is Ptr { - return isMarkedInternal(schema) && (schema as Ptr)?.type === 'ptr'; + return (schema as Ptr)?.type === 'ptr' && isMarkedInternal(schema); } /** @@ -1704,44 +1704,44 @@ export function isPtr(schema: unknown): schema is Ptr { * isAtomic(d.u32) // false */ export function isAtomic(schema: unknown): schema is Atomic { - return isMarkedInternal(schema) && (schema as Atomic)?.type === 'atomic'; + return (schema as Atomic)?.type === 'atomic' && isMarkedInternal(schema); } export function isAlignAttrib(value: unknown): value is Align { - return isMarkedInternal(value) && (value as Align)?.type === '@align'; + return (value as Align)?.type === '@align' && isMarkedInternal(value); } export function isSizeAttrib(value: unknown): value is Size { - return isMarkedInternal(value) && (value as Size)?.type === '@size'; + return (value as Size)?.type === '@size' && isMarkedInternal(value); } export function isLocationAttrib(value: unknown): value is Location { - return isMarkedInternal(value) && (value as Location)?.type === '@location'; + return (value as Location)?.type === '@location' && isMarkedInternal(value); } export function isInterpolateAttrib( value: unknown, ): value is Interpolate { - return isMarkedInternal(value) && (value as Interpolate)?.type === '@interpolate'; + return (value as Interpolate)?.type === '@interpolate' && isMarkedInternal(value); } export function isBuiltinAttrib(value: unknown): value is Builtin { - return isMarkedInternal(value) && (value as Builtin)?.type === '@builtin'; + return (value as Builtin)?.type === '@builtin' && isMarkedInternal(value); } export function isInvariantAttrib(value: unknown): value is Invariant { - return isMarkedInternal(value) && (value as Invariant)?.type === '@invariant'; + return (value as Invariant)?.type === '@invariant' && isMarkedInternal(value); } export function isDecorated(value: unknown): value is Decorated { - return isMarkedInternal(value) && (value as Decorated)?.type === 'decorated'; + return (value as Decorated)?.type === 'decorated' && isMarkedInternal(value); } export function isAbstractFloat(value: unknown): value is AbstractFloat { - return isMarkedInternal(value) && (value as AbstractFloat).type === 'abstractFloat'; + return (value as AbstractFloat)?.type === 'abstractFloat' && isMarkedInternal(value); } export function isAbstractInt(value: unknown): value is AbstractInt { - return isMarkedInternal(value) && (value as AbstractInt).type === 'abstractInt'; + return (value as AbstractInt)?.type === 'abstractInt' && isMarkedInternal(value); } export function isAbstract(value: unknown): value is AbstractFloat | AbstractInt { @@ -1753,11 +1753,11 @@ export function isConcrete(value: unknown): boolean { } export function isVoid(value: unknown): value is Void { - return isMarkedInternal(value) && (value as Void).type === 'void'; + return (value as Void)?.type === 'void' && isMarkedInternal(value); } export function isBool(value: unknown): value is Bool { - return isMarkedInternal(value) && (value as Bool).type === 'bool'; + return (value as Bool)?.type === 'bool' && isMarkedInternal(value); } export function isNumericSchema( @@ -1766,13 +1766,13 @@ export function isNumericSchema( const type = (schema as BaseData)?.type; return ( - isMarkedInternal(schema) && (type === 'abstractInt' || type === 'abstractFloat' || type === 'f32' || type === 'f16' || type === 'i32' || - type === 'u32') + type === 'u32') && + isMarkedInternal(schema) ); } @@ -1780,8 +1780,8 @@ export function isHalfPrecisionSchema(schema: unknown): schema is F16 | Vec2h | const type = (schema as BaseData)?.type; return ( - isMarkedInternal(schema) && - (type === 'f16' || type === 'vec2h' || type === 'vec3h' || type === 'vec4h') + (type === 'f16' || type === 'vec2h' || type === 'vec3h' || type === 'vec4h') && + isMarkedInternal(schema) ); } diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index c183d04cdd..b6dad8f68f 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -66,7 +66,13 @@ import { isTgpuFn } from './core/function/tgpuFn.ts'; import type { IOData } from './core/function/fnTypes.ts'; import { AutoStruct } from './data/autoStruct.ts'; import { EntryInputRouter } from './core/function/entryInputRouter.ts'; -import { validateIdentifier, sanitizePrimer, bannedTokens } from './nameUtils.ts'; +import { + validateIdentifier, + sanitizePrimer, + bannedTokens, + builtins, + templateEnumerants, +} from './nameUtils.ts'; import { minify } from './minify.ts'; /** @@ -155,6 +161,7 @@ class ItemStateStackImpl implements ItemStateStack { argAccess, returnType, externalMap, + externalSnippets: new Map(), reportedReturnTypes: new Set(), placeholderForVariable: new Map(), modifiedVariables: new Set(), @@ -220,12 +227,18 @@ class ItemStateStackImpl implements ItemStateStack { return access(); } + const cached = layer.externalSnippets.get(id); + if (cached) { + return cached; + } if (Object.hasOwn(layer.externalMap, id)) { const external = layer.externalMap[id]; if (isNamable(external) && getName(external) === undefined) { setName(external, id.replaceAll('.', '_')); } - return coerceToSnippet(external); + const snippet = coerceToSnippet(external); + layer.externalSnippets.set(id, snippet); + return snippet; } return undefined; @@ -491,6 +504,9 @@ export class ResolutionCtxImpl implements ResolutionCtx { isIdentifierTaken(name: string, scope: 'global' | 'block'): boolean { return ( + bannedTokens.has(name) || + builtins.has(name) || + templateEnumerants.has(name) || this.#namespaceInternal.takenGlobalIdentifiers.has(name) || (scope === 'block' ? this._itemStateStack.isIdentifierTakenLocally(name) @@ -911,6 +927,18 @@ export class ResolutionCtxImpl implements ResolutionCtx { return maybeEventual; } + /** + * Whether every slot in the map currently has a value equal to the one recorded in it. + */ + #slotsHaveValues(slotToValueMap: Map, unknown>): boolean { + for (const [slot, expectedValue] of slotToValueMap) { + if (!slot.areEqual(this._itemStateStack.readSlot(slot), expectedValue)) { + return false; + } + } + return true; + } + _getOrCompute(lazy: TgpuLazy): T { // All memoized versions of `lazy` const instances = this.#namespaceInternal.memoizedLazy.get(lazy) ?? []; @@ -919,13 +947,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { try { for (const instance of instances) { - const slotValuePairs = [...instance.slotToValueMap.entries()]; - - if ( - slotValuePairs.every(([slot, expectedValue]) => - slot.areEqual(this._itemStateStack.readSlot(slot), expectedValue), - ) - ) { + if (this.#slotsHaveValues(instance.slotToValueMap)) { return instance.result as T; } } @@ -972,13 +994,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { try { for (const instance of instances) { - const slotValuePairs = [...instance.slotToValueMap.entries()]; - - if ( - slotValuePairs.every(([slot, expectedValue]) => - slot.areEqual(this._itemStateStack.readSlot(slot), expectedValue), - ) - ) { + if (this.#slotsHaveValues(instance.slotToValueMap)) { return instance.result; } } diff --git a/packages/typegpu/src/tgsl/conversion.ts b/packages/typegpu/src/tgsl/conversion.ts index 1a84ac3d7d..5a0bf42cbc 100644 --- a/packages/typegpu/src/tgsl/conversion.ts +++ b/packages/typegpu/src/tgsl/conversion.ts @@ -145,10 +145,37 @@ function getImplicitConversionRank(src: BaseData, dest: BaseData): ConversionRan return INFINITE_RANK; } +/** + * Conversion ranks only depend on immutable facts about the two schemas (their type tags, + * identity, inner types and pointer flags), so they're computed once per pair. Callers must + * treat the returned info as read-only. + */ +const autoRankCache = new WeakMap>(); +const implicitRankCache = new WeakMap>(); + function getConversionRank( src: BaseData, dest: BaseData, allowImplicit: boolean, +): ConversionRankInfo { + const cache = allowImplicit ? implicitRankCache : autoRankCache; + let byDest = cache.get(src); + if (!byDest) { + byDest = new WeakMap(); + cache.set(src, byDest); + } + let rank = byDest.get(dest); + if (!rank) { + rank = computeConversionRank(src, dest, allowImplicit); + byDest.set(dest, rank); + } + return rank; +} + +function computeConversionRank( + src: BaseData, + dest: BaseData, + allowImplicit: boolean, ): ConversionRankInfo { const autoRank = getAutoConversionRank(src, dest); if (autoRank.rank < Number.POSITIVE_INFINITY) { diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 26f66067e4..055644b27f 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -807,7 +807,7 @@ export class WgslGenerator implements ShaderGenerator { typeof expression[1] === 'string' ? numericLiteralToSnippet(parseNumericString(expression[1])) : numericLiteralToSnippet(expression[1]); - invariant(type, `Expected ${stringifyNode(expression)} to be valid numeric literal`); + invariant(type, () => `Expected ${stringifyNode(expression)} to be valid numeric literal`); return type; } @@ -1266,17 +1266,17 @@ The generated shader will evaluate them in declaration order.`, let body = this._block(options.body, /* allowInlining */ false); const scope = this.ctx.topFunctionScope; invariant(scope, 'Expected function scope to be present'); - const replacements = Object.fromEntries( - [...scope.placeholderForVariable.entries()].map(([variable, placeholder]) => [ - placeholder, - scope.modifiedVariables.has(variable) ? 'var' : 'let', - ]), - ); - if (Object.keys(replacements).length > 0) { - const regex = new RegExp(Object.keys(replacements).join('|'), 'gi'); + if (scope.placeholderForVariable.size > 0) { + const keywordForPlaceholder = new Map(); + for (const [variable, placeholder] of scope.placeholderForVariable) { + keywordForPlaceholder.set( + placeholder, + scope.modifiedVariables.has(variable) ? 'var' : 'let', + ); + } body.code = body.code.replace( - regex, - (match) => replacements[match as keyof typeof replacements] ?? '#ERR', + /#VAR_\d+#/g, + (match) => keywordForPlaceholder.get(match) ?? '#ERR', ); } diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts index 3c66e1031f..b1d2293a77 100644 --- a/packages/typegpu/src/types.ts +++ b/packages/typegpu/src/types.ts @@ -119,6 +119,8 @@ export type FunctionScopeLayer = { functionType: 'normal' | 'compute' | 'vertex' | 'fragment'; argAccess: Record; externalMap: Record; + /** Snippets of externals already looked up in this function, so each is only coerced once */ + externalSnippets: Map; /** * The return type of the function. If undefined, the type should be inferred * from the implementation (relevant for shellless functions).