diff --git a/apps/typegpu-docs/src/content/docs/reference/shader-generation.mdx b/apps/typegpu-docs/src/content/docs/reference/shader-generation.mdx index ac1a2267f..209bbd6ba 100644 --- a/apps/typegpu-docs/src/content/docs/reference/shader-generation.mdx +++ b/apps/typegpu-docs/src/content/docs/reference/shader-generation.mdx @@ -200,6 +200,7 @@ There are essentially three types of origins: include accessing the position of an existing `Boid` instance or modifying the position of an existing `Boid` instance. These include `'uniform'`, `'mutable'`, `'readonly'`, `'workgroup'`, `'private'`, `'function'`, `'handle'`, `'constant-ref'` and `'this-function'`. - `'uniform'`, `'mutable'`, `'readonly'`, `'workgroup'`, `'private'`, `'function'`, `'handle'` all reflect address spaces that values can belong to, and we use them to determine what kind of pointer type they are. - - `'constant-ref'` is a reference to a value stored in a `tgpu.const`. They're different from `constant`s, as we know that even if they're referential (non-primitive), the developer cannot mutate them. + - `'constant-tgpu-const-ref'` is a reference to a value stored in a `tgpu.const`. They're different from `constant`s, as we know that even if they're referential (non-primitive), the developer cannot mutate them. + - `'runtime-tgpu-const-ref'` is an expression that still references a `tgpu.const` value, but isn't considered constant in the eyes of WGSL. - `'this-function'` lets us track whether values originates from the function we're currently generating, or the function that called us. - **Argument Origins**: This group is dedicated to exactly one origin: 'argument'. It represents values that are passed as arguments to functions. diff --git a/packages/typegpu/src/core/constant/tgpuConstant.ts b/packages/typegpu/src/core/constant/tgpuConstant.ts index a6cea2228..915b635f4 100644 --- a/packages/typegpu/src/core/constant/tgpuConstant.ts +++ b/packages/typegpu/src/core/constant/tgpuConstant.ts @@ -98,7 +98,9 @@ class TgpuConstImpl return snip( id, this.dataType, - isNaturallyEphemeral(this.dataType) ? 'constant' : 'constant-ref', + isNaturallyEphemeral(this.dataType) + ? 'constant' + : 'constant-tgpu-const-ref', ); } @@ -115,7 +117,9 @@ class TgpuConstImpl return snip( this, dataType, - isNaturallyEphemeral(dataType) ? 'constant' : 'constant-ref', + isNaturallyEphemeral(dataType) + ? 'constant' + : 'constant-tgpu-const-ref', ); }, [$resolve]: (ctx) => ctx.resolve(this), diff --git a/packages/typegpu/src/data/snippet.ts b/packages/typegpu/src/data/snippet.ts index ceee72370..b67fc5006 100644 --- a/packages/typegpu/src/data/snippet.ts +++ b/packages/typegpu/src/data/snippet.ts @@ -24,7 +24,10 @@ export type Origin = // note that this doesn't automatically mean the value can be stored in a `const` // variable, more so that it's valid to do so in WGSL (but not necessarily safe to do in JS shaders) | 'constant' - | 'constant-ref'; + // don't even get me started on these. They're references to non-primitive values that originate + // from a tgpu.const(...).$ call. + | 'constant-tgpu-const-ref' /* turns into a `const` when assigned to a variable */ + | 'runtime-tgpu-const-ref' /* turns into a `let` when assigned to a variable */; export function isEphemeralOrigin(space: Origin) { return space === 'runtime' || space === 'constant' || space === 'argument'; diff --git a/packages/typegpu/src/tgsl/generationHelpers.ts b/packages/typegpu/src/tgsl/generationHelpers.ts index 159926719..f0c8e94c0 100644 --- a/packages/typegpu/src/tgsl/generationHelpers.ts +++ b/packages/typegpu/src/tgsl/generationHelpers.ts @@ -20,6 +20,7 @@ import { import { isEphemeralSnippet, isSnippet, + type Origin, snip, type Snippet, } from '../data/snippet.ts'; @@ -156,10 +157,7 @@ export function accessProp( target: Snippet, propName: string, ): Snippet | undefined { - if ( - infixKinds.includes(target.dataType.type) && - propName in infixOperators - ) { + if (infixKinds.includes(target.dataType.type) && propName in infixOperators) { return snip( new InfixDispatch( propName, @@ -174,11 +172,7 @@ export function accessProp( if (isWgslArray(target.dataType) && propName === 'length') { if (target.dataType.elementCount === 0) { // Dynamically-sized array - return snip( - stitch`arrayLength(&${target})`, - u32, - /* origin */ 'runtime', - ); + return snip(stitch`arrayLength(&${target})`, u32, /* origin */ 'runtime'); } return snip( @@ -208,10 +202,10 @@ export function accessProp( propType, /* origin */ target.origin === 'argument' ? 'argument' - : !isEphemeralSnippet(target) && - !isNaturallyEphemeral(propType) + : !isEphemeralSnippet(target) && !isNaturallyEphemeral(propType) ? target.origin - : target.origin === 'constant' || target.origin === 'constant-ref' + : target.origin === 'constant' || + target.origin === 'constant-tgpu-const-ref' ? 'constant' : 'runtime', ); @@ -231,11 +225,7 @@ export function accessProp( } const propLength = propName.length; - if ( - isVec(target.dataType) && - propLength >= 1 && - propLength <= 4 - ) { + if (isVec(target.dataType) && propLength >= 1 && propLength <= 4) { const swizzleTypeChar = target.dataType.type.includes('bool') ? 'b' : (target.dataType.type[4] as SwizzleableType); @@ -255,7 +245,7 @@ export function accessProp( /* origin */ target.origin === 'argument' && propLength === 1 ? 'argument' : target.origin === 'constant' || - target.origin === 'constant-ref' + target.origin === 'constant-tgpu-const-ref' ? 'constant' : 'runtime', ); @@ -282,6 +272,28 @@ export function accessIndex( // array if (isWgslArray(target.dataType) || isDisarray(target.dataType)) { const elementType = target.dataType.elementType as AnyData; + const isTargetEphemeral = isEphemeralSnippet(target); + const isElementNatEph = isNaturallyEphemeral(elementType); + + let origin: Origin; + + if ( + target.origin === 'constant-tgpu-const-ref' && + index.origin === 'constant' + ) { + origin = isElementNatEph ? 'constant' : 'constant-tgpu-const-ref'; + } else if ( + target.origin === 'constant-tgpu-const-ref' || + target.origin === 'runtime-tgpu-const-ref' + ) { + origin = isElementNatEph ? 'runtime' : 'runtime-tgpu-const-ref'; + } else if (!isTargetEphemeral && !isElementNatEph) { + origin = target.origin; + } else if (target.origin === 'constant' && index.origin === 'constant') { + origin = 'constant'; + } else { + origin = 'runtime'; + } return snip( isKnownAtComptime(target) && isKnownAtComptime(index) @@ -289,12 +301,7 @@ export function accessIndex( ? (target.value as any)[index.value as number] : stitch`${target}[${index}]`, elementType, - /* origin */ !isEphemeralSnippet(target) && - !isNaturallyEphemeral(elementType) - ? target.origin - : target.origin === 'constant' || target.origin === 'constant-ref' - ? 'constant' - : 'runtime', + /* origin */ origin, ); } @@ -307,7 +314,7 @@ export function accessIndex( : stitch`${target}[${index}]`, target.dataType.primitive, /* origin */ target.origin === 'constant' || - target.origin === 'constant-ref' + target.origin === 'constant-tgpu-const-ref' ? 'constant' : 'runtime', ); diff --git a/packages/typegpu/src/tgsl/shellless.ts b/packages/typegpu/src/tgsl/shellless.ts index 429309653..ff7c877b5 100644 --- a/packages/typegpu/src/tgsl/shellless.ts +++ b/packages/typegpu/src/tgsl/shellless.ts @@ -67,7 +67,10 @@ export class ShelllessRepository { let type = concretize(s.dataType as AnyData); - if (s.origin === 'constant-ref') { + if ( + s.origin === 'constant-tgpu-const-ref' || + s.origin === 'runtime-tgpu-const-ref' + ) { // biome-ignore lint/style/noNonNullAssertion: it's there const ctx = getResolutionCtx()!; throw new Error( diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 7149274e7..ac98d7bbc 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -172,10 +172,13 @@ ${this.ctx.pre}}`; const naturallyEphemeral = wgsl.isNaturallyEphemeral(dataType); let varOrigin: Origin; - if (origin === 'constant-ref') { + if ( + origin === 'constant-tgpu-const-ref' || + origin === 'runtime-tgpu-const-ref' + ) { // Even types that aren't naturally referential (like vectors or structs) should // be treated as constant references when assigned to a const. - varOrigin = 'constant-ref'; + varOrigin = origin; } else if (origin === 'argument') { if (naturallyEphemeral) { varOrigin = 'runtime'; @@ -288,7 +291,9 @@ ${this.ctx.pre}}`; if (exprType === NODE.assignmentExpr) { if ( - convLhs.origin === 'constant' || convLhs.origin === 'constant-ref' + convLhs.origin === 'constant' || + convLhs.origin === 'constant-tgpu-const-ref' || + convLhs.origin === 'runtime-tgpu-const-ref' ) { throw new WgslTypeError( `'${lhsStr} = ${rhsStr}' is invalid, because ${lhsStr} is a constant.`, @@ -875,8 +880,10 @@ ${this.ctx.pre}else ${alternate}`; ); } - if (eq.origin === 'constant-ref') { + if (eq.origin === 'constant-tgpu-const-ref') { varType = 'const'; + } else if (eq.origin === 'runtime-tgpu-const-ref') { + varType = 'let'; } else { varType = 'let'; if (!wgsl.isPtr(dataType)) { diff --git a/packages/typegpu/tests/constant.test.ts b/packages/typegpu/tests/constant.test.ts index f27084d29..392ebcc11 100644 --- a/packages/typegpu/tests/constant.test.ts +++ b/packages/typegpu/tests/constant.test.ts @@ -91,4 +91,34 @@ describe('tgpu.const', () => { `[TypeError: Cannot assign to read only property 'pos' of object '#']`, ); }); + + it('looses its `constant` origin when indexing with runtime value', () => { + const positions = tgpu.const(d.arrayOf(d.vec3f, 3), [ + d.vec3f(0), + d.vec3f(1), + d.vec3f(2), + ]); + + const foo = (idx: number) => { + 'use gpu'; + const pos = positions.$[idx]; + }; + + const main = () => { + 'use gpu'; + foo(0); + }; + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "const positions: array = array(vec3f(), vec3f(1), vec3f(2)); + + fn foo(idx: i32) { + let pos = positions[idx]; + } + + fn main() { + foo(0i); + }" + `); + }); });