typegpu 0.12.7 → 0.12.8

package/tgsl/wgslGenerator.js

@@ -4,7 +4,7 @@ import { arrayOf } from "../data/array.js";
 import { UnknownData, unptr } from "../data/dataTypes.js";
 import { bool, i32, u32 } from "../data/numeric.js";
 import { vec2u, vec3u, vec4u } from "../data/vector.js";
-import { fallthroughCopyOrigin, isStoredInMemory, snip, } from "../data/snippet.js";
+import { fallthroughCopyOrigin, isMutableOrigin, isStoredInMemory, snip, } from "../data/snippet.js";
 import * as wgsl from "../data/wgslTypes.js";
 import { invariant, ResolutionError, WgslForbiddenStatementError, WgslTypeError, } from "../errors.js";
 import { getName } from "../shared/meta.js";
@@ -13,7 +13,7 @@ import { safeStringify } from "../shared/stringify.js";
 import { pow } from "../std/numeric.js";
 import { add, div, mul, neg, sub } from "../std/operators.js";
 import { eq, ne, lt, le, gt, ge, not } from "../std/boolean.js";
-import { isConstant, isGPUCallable, isKnownAtComptime, } from "../types.js";
+import { asComptime, isConstant, isGPUCallable, } from "../types.js";
 import { convertStructValues, convertToCommonType, tryConvertSnippet } from "./conversion.js";
 import { ArrayExpression, coerceToSnippet, concretize, numericLiteralToSnippet, } from "./generationHelpers.js";
 import { accessIndex } from "./accessIndex.js";
@@ -29,6 +29,7 @@ import { mathToStd, supportedLogOps } from "./jsPolyfills.js";
 import * as forOfUtils from "./forOfUtils.js";
 import { isTgpuRange } from "../std/range.js";
 import { stringifyNode, stringifyObjectProperty } from "../shared/tseynit.js";
+import { setSourceNode, stringifySnippet } from "./stringifySnippet.js";
 import { getAttributesString } from "../data/attributes.js";
 import { validSelectBranchTypes } from "../std/boolean.js";
 import { isInfixDispatch } from "./infixDispatch.js";
@@ -127,12 +128,13 @@ const unaryOpCodeToCodegen = {
         if (argExpr === undefined) {
             throw new Error('The unary operator `!` expects 1 argument, but 0 were provided.');
         }
-        if (isKnownAtComptime(argExpr)) {
-            return snip(!argExpr.value, bool, 'constant', false);
+        const knownArg = asComptime(argExpr);
+        if (knownArg) {
+            return snip(!knownArg.value, bool, 'constant', false);
         }
         const argStr = ctx.resolveSnippet(argExpr).value;
         if (!wgsl.isBool(argExpr.dataType)) {
-            throw new WgslTypeError(`Unary operator ! requires boolean operand. Got ${String(argExpr.dataType)}.${wgsl.isVecBool(argExpr.dataType)
+            throw new WgslTypeError(`Unary operator ! requires boolean operand. Got '${stringifySnippet(argExpr)}' of type ${String(argExpr.dataType)}.${wgsl.isVecBool(argExpr.dataType)
                 ? ` For component-wise negation, use 'std.${not.toString()}'.`
                 : ''}`);
         }
@@ -279,25 +281,41 @@ export class WgslGenerator {
         return res;
     }
     _emitSwitchStatement(discriminantExpr, groupedCaseExprs) {
-        if (groupedCaseExprs.flatMap(([tests]) => tests).every((test) => test.value !== 'default')) {
-            // default clause is required in WGSL
-            groupedCaseExprs.push([[switchDefault], []]);
-        }
-        this.ctx.indent();
         const cases = groupedCaseExprs.map(([tests, consequent]) => {
-            const resolvedTests = tests
-                .map((test) => (test.value === 'default' ? 'default' : this.ctx.resolveSnippet(test).value))
-                .join(', ');
-            // Purely cosmetic break pruning (it is legal, but there's no need to generate it).
+            // It is legal to generate a break at the end of a consequent, but there is no need to.
+            // Also, removing this lets us simplify more pruned switch statements to blocks.
             const last = consequent.at(-1);
             if (last && /^\s*break\s*;\s*$/.test(last.code) && last.endsWithControlFlow === 'break') {
                 consequent.pop();
             }
             const resolvedConsequent = consequent.map((s) => s.code).join('\n');
-            return stitch `${this.ctx.pre}case ${resolvedTests}: {\n${resolvedConsequent}\n${this.ctx.pre}}`;
+            return { tests, resolvedConsequent };
+        });
+        if (cases.flatMap(({ tests }) => tests).every((test) => test !== switchDefault)) {
+            // default clause is required in WGSL
+            cases.push({ tests: [switchDefault], resolvedConsequent: '' });
+        }
+        // If the switch was pruned, and it has no breaks, we can simplify it to a block.
+        const onlyCase = cases.length === 1 ? cases.at(0) : undefined;
+        if (asComptime(discriminantExpr) !== undefined &&
+            onlyCase &&
+            onlyCase.tests.length === 1 &&
+            onlyCase.tests.at(0) === switchDefault) {
+            const { resolvedConsequent } = onlyCase;
+            if (!/\bbreak\b/.test(resolvedConsequent)) {
+                return `${this.ctx.pre}{\n${this.ctx.getDedented(resolvedConsequent)}\n${this.ctx.pre}}`;
+            }
+        }
+        this.ctx.indent();
+        const resolvedCases = cases.map(({ tests, resolvedConsequent }) => {
+            const resolvedTests = tests
+                .map((test) => (test.value === 'default' ? 'default' : this.ctx.resolveSnippet(test).value))
+                .join(', ');
+            return `${this.ctx.pre}case ${resolvedTests}: {\n${resolvedConsequent}\n${this.ctx.pre}}`;
         });
         this.ctx.dedent();
-        return stitch `${this.ctx.pre}switch ${discriminantExpr} {\n${cases.join('\n')}\n${this.ctx.pre}}`;
+        const body = resolvedCases.join('\n');
+        return stitch `${this.ctx.pre}switch ${discriminantExpr} {\n${body}\n${this.ctx.pre}}`;
     }
     _callShellless(callee, args) {
         const isGeneric = isGenericFn(callee);
@@ -335,13 +353,17 @@ export class WgslGenerator {
                 // convert the result.
                 return result;
             }
-            return tryConvertSnippet(this.ctx, result, expectedType);
+            return setSourceNode(tryConvertSnippet(this.ctx, result, expectedType), expression);
         }
         finally {
             this.ctx.expectedType = prevExpectedType;
         }
     }
     _expression(expression) {
+        // Remembering which JS expression produced the snippet, for more descriptive errors
+        return setSourceNode(this._expressionImpl(expression), expression);
+    }
+    _expressionImpl(expression) {
         if (isId(expression)) {
             return this._identifier(extractId(expression));
         }
@@ -352,21 +374,23 @@ export class WgslGenerator {
             const [_, lhs, op, rhs] = expression;
             const lhsExpr = this._expression(lhs);
             // Short Circuit Evaluation
-            if (isKnownAtComptime(lhsExpr)) {
+            const knownLhs = asComptime(lhsExpr);
+            if (knownLhs) {
                 const castToBool = wgsl.isBool(this.ctx.expectedType);
-                const evalRhs = op === '&&' ? lhsExpr.value : !lhsExpr.value;
+                const evalRhs = op === '&&' ? knownLhs.value : !knownLhs.value;
                 if (!evalRhs) {
                     return castToBool
                         ? snip(op === '||', bool, 'constant', false)
-                        : coerceToSnippet(lhsExpr.value);
+                        : coerceToSnippet(knownLhs.value);
                 }
                 const rhsExpr = this._expression(rhs);
-                if (isKnownAtComptime(rhsExpr)) {
-                    const rhsSnippet = coerceToSnippet(rhsExpr.value);
+                const knownRhs = asComptime(rhsExpr);
+                if (knownRhs) {
+                    const rhsSnippet = coerceToSnippet(knownRhs.value);
                     return castToBool ? tryConvertSnippet(this.ctx, rhsSnippet, bool, false) : rhsSnippet;
                 }
                 if (rhsExpr.dataType === UnknownData) {
-                    throw new WgslTypeError(`Right-hand side of '${op}' is of unknown type`);
+                    throw new WgslTypeError(`Right-hand side '${stringifySnippet(rhsExpr)}' of '${op}' is of unknown type.`);
                 }
                 // we can skip lhs
                 return castToBool ? tryConvertSnippet(this.ctx, rhsExpr, bool, false) : rhsExpr;
@@ -374,13 +398,13 @@ export class WgslGenerator {
             const rhsExpr = this._expression(rhs);
             // they are not known at comptime
             if (lhsExpr.dataType === UnknownData) {
-                throw new WgslTypeError(`Left-hand side of '${op}' is of unknown type`);
+                throw new WgslTypeError(`Left-hand side '${stringifySnippet(lhsExpr)}' of '${op}' is of unknown type.`);
             }
-            if (!isKnownAtComptime(rhsExpr) && rhsExpr.dataType === UnknownData) {
-                throw new WgslTypeError(`Right-hand side of '${op}' is of unknown type`);
+            if (!asComptime(rhsExpr) && rhsExpr.dataType === UnknownData) {
+                throw new WgslTypeError(`Right-hand side '${stringifySnippet(rhsExpr)}' of '${op}' is of unknown type.`);
             }
             if (!wgsl.isBool(lhsExpr.dataType) || !wgsl.isBool(rhsExpr.dataType)) {
-                throw new WgslTypeError(`Logical expression '${op}' requires boolean operands. Got '${String(lhsExpr.dataType)}' and '${String(rhsExpr.dataType)}'.`);
+                throw new WgslTypeError(`Logical expression '${stringifyNode(expression)}' requires boolean operands. Got '${String(lhsExpr.dataType)}' and '${String(rhsExpr.dataType)}'.`);
             }
             const lhsStr = this.ctx.resolveSnippet(lhsExpr).value;
             const rhsStr = this.ctx.resolveSnippet(rhsExpr).value;
@@ -393,7 +417,7 @@ export class WgslGenerator {
             const lhsExpr = this._expression(lhs);
             const rhsExpr = this._expression(rhs);
             if (rhsExpr.value instanceof RefOperator) {
-                throw new WgslTypeError(stitch `Cannot assign a ref to an existing variable '${stringifyNode(lhs)}', define a new variable instead.`);
+                throw new WgslTypeError(`Cannot assign a ref to an existing variable '${stringifyNode(lhs)}', define a new variable instead.`);
             }
             if (op === '==') {
                 throw new Error('Please use the === operator instead of ==');
@@ -402,9 +426,11 @@ export class WgslGenerator {
                 throw new Error('Please use the !== operator instead of !=');
             }
             const stdBinaryRelationalOp = binaryRelationalOpToStdMap[op];
-            if (stdBinaryRelationalOp && isKnownAtComptime(lhsExpr) && isKnownAtComptime(rhsExpr)) {
-                const left = lhsExpr.value;
-                const right = rhsExpr.value;
+            const knownLhs = stdBinaryRelationalOp && asComptime(lhsExpr);
+            const knownRhs = knownLhs && asComptime(rhsExpr);
+            if (stdBinaryRelationalOp && knownLhs && knownRhs) {
+                const left = knownLhs.value;
+                const right = knownRhs.value;
                 switch (op) {
                     case '===':
                         return snip(left === right, bool, 'constant', false);
@@ -413,7 +439,7 @@ export class WgslGenerator {
                 }
                 if (typeof left !== 'number' || typeof right !== 'number') {
                     const bothVectors = wgsl.isVec(lhsExpr.dataType) && wgsl.isVec(rhsExpr.dataType);
-                    throw new WgslTypeError(`Comparison '${op}' requires numeric operands.${bothVectors
+                    throw new WgslTypeError(`Comparison '${stringifyNode(expression)}' requires numeric operands.${bothVectors
                         ? ` For component-wise comparison, use 'std.${stdBinaryRelationalOp}'.`
                         : ''}`);
                 }
@@ -429,10 +455,10 @@ export class WgslGenerator {
                 }
             }
             if (lhsExpr.dataType === UnknownData) {
-                throw new WgslTypeError(`Left-hand side of '${op}' is of unknown type`);
+                throw new WgslTypeError(`Left-hand side '${stringifySnippet(lhsExpr)}' of '${op}' is of unknown type.`);
             }
             if (rhsExpr.dataType === UnknownData) {
-                throw new WgslTypeError(`Right-hand side of '${op}' is of unknown type`);
+                throw new WgslTypeError(`Right-hand side '${stringifySnippet(rhsExpr)}' of '${op}' is of unknown type.`);
             }
             const codegen = binaryOpCodeToCodegen[op];
             if (codegen) {
@@ -488,7 +514,7 @@ export class WgslGenerator {
                     (equalityCheck && wgsl.isBool(convLhs.dataType) && wgsl.isBool(convRhs.dataType));
                 if (!correctOperandTypes) {
                     const bothVectors = wgsl.isVec(convLhs.dataType) && wgsl.isVec(convRhs.dataType);
-                    throw new WgslTypeError(`Comparison '${op}' requires numeric${equalityCheck ? ' or boolean' : ''} operands. Got '${String(convLhs.dataType)}' and '${String(convRhs.dataType)}'.${bothVectors
+                    throw new WgslTypeError(`Comparison '${stringifyNode(expression)}' requires numeric${equalityCheck ? ' or boolean' : ''} operands. Got '${String(convLhs.dataType)}' and '${String(convRhs.dataType)}'.${bothVectors
                         ? ` For component-wise comparison, use 'std.${stdBinaryRelationalOp}'.`
                         : ''}`);
                 }
@@ -662,7 +688,7 @@ export class WgslGenerator {
             if (maybeConsoleMethod) {
                 throw new Error(`Unsupported console functionality 'console.${maybeConsoleMethod}()'.`);
             }
-            throw new Error(`Function '${getName(callee.value) ?? String(callee.value)}' is not marked with the 'use gpu' directive and cannot be used in a shader`);
+            throw new Error(`Function '${getName(callee.value) ?? stringifySnippet(callee)}' is not marked with the 'use gpu' directive and cannot be used in a shader`);
         }
         if (expression[0] === NODE.objectExpr) {
             // Normalize to `objectProperty[]`
@@ -694,7 +720,7 @@ export class WgslGenerator {
                         // Generating the expression and inferring the type instead
                         expr = this._expression(value);
                         if (expr.dataType === UnknownData) {
-                            throw new WgslTypeError(stitch `Property ${key} in object literal has a value of unknown type: '${expr}'`);
+                            throw new WgslTypeError(`Value '${stringifySnippet(expr)}' of property '${key}' in object literal is of unknown type.`);
                         }
                         // Taking care of abstract numerics and implicit pointers
                         accessed = structType.provideProp(key, unptr(concretize(expr.dataType)));
@@ -789,8 +815,11 @@ The generated shader will evaluate them in declaration order.`);
             // ternary operator
             const [_, testNode, consequentNode, alternativeNode] = expression;
             const test = this._expression(testNode);
-            if (isKnownAtComptime(test)) {
-                return test.value ? this._expression(consequentNode) : this._expression(alternativeNode);
+            const knownTest = asComptime(test);
+            if (knownTest) {
+                return knownTest.value
+                    ? this._expression(consequentNode)
+                    : this._expression(alternativeNode);
             }
             else {
                 const convertedTest = tryConvertSnippet(this.ctx, test, bool, false);
@@ -1055,9 +1084,11 @@ Try 'return ${typeStr}(${str});' instead.
         const eq = this._expression(eqNode);
         if (eq.value instanceof RefOperator) {
             // We're assigning a newly created `d.ref()`
-            if (eq.dataType !== UnknownData) {
-                throw new WgslTypeError(`Cannot store d.ref() in a variable if it references another value. Copy the value passed into d.ref() instead.`);
+            if (eq.value.addressable) {
+                const refStr = stringifyNode(eqNode);
+                throw new WgslTypeError(`'const ${rawId} = ${refStr}' is invalid, cannot store ${refStr} in a variable if it references another value. Copy the value passed into d.ref() instead, or inline ${refStr}.`);
             }
+            // Unwrapping the ref, and storing the value it was created from in a new variable
             const refSnippet = eq.value.snippet;
             const varName = this.refVariable(rawId, concretize(refSnippet.dataType));
             return {
@@ -1287,7 +1318,8 @@ ${this.ctx.pre}else ${alternate}`,
                 const originalLoopVarName = extractId(loopVar[1]);
                 const blockified = blockifySingleStatement(body);
                 if (shouldUnroll) {
-                    if (!isKnownAtComptime(range.end)) {
+                    const knownEnd = asComptime(range.end);
+                    if (!knownEnd) {
                         throw new Error('Cannot unroll loop. Length of iterable is unknown at comptime.');
                     }
                     const { value } = iterableSnippet;
@@ -1295,7 +1327,7 @@ ${this.ctx.pre}else ${alternate}`,
                         ? value.map((i) => coerceToSnippet(i))
                         : value instanceof ArrayExpression
                             ? value.elements
-                            : Array.from({ length: range.end.value }, (_, i) => forOfUtils.getElementSnippet(iterableSnippet, snip(i, u32, 'constant')));
+                            : Array.from({ length: knownEnd.value }, (_, i) => forOfUtils.getElementSnippet(iterableSnippet, snip(i, u32, 'constant')));
                     if (elements.length === 0) {
                         return { code: '', definesInNearestScope: false };
                     }
@@ -1373,8 +1405,10 @@ ${this.ctx.pre}else ${alternate}`,
             });
             // comptime folding
             let matchedCaseWasNotLast = false;
-            if ([discriminantExpr, ...caseExprs.map(([test]) => test)].every(isKnownAtComptime)) {
-                let matchedCaseIndex = caseExprs.findIndex(([test]) => test !== switchDefault && test.value === discriminantExpr.value);
+            const knownDiscriminant = asComptime(discriminantExpr);
+            const knownTests = knownDiscriminant && caseExprs.map(([test]) => asComptime(test));
+            if (knownDiscriminant && knownTests?.every((test) => test !== undefined)) {
+                let matchedCaseIndex = knownTests.findIndex((test) => test !== switchDefault && test?.value === knownDiscriminant.value);
                 if (matchedCaseIndex === -1) {
                     matchedCaseIndex = caseExprs.findIndex(([test]) => test === switchDefault);
                 }
@@ -1517,8 +1551,8 @@ ${stringifyNode(statement)}`);
         if (!computed) {
             return property[0];
         }
-        const key = this._expression(property[0]);
-        if (!isKnownAtComptime(key)) {
+        const key = asComptime(this._expression(property[0]));
+        if (!key) {
             throw new WgslTypeError(`Computed object property key '${stringifyObjectProperty(property)}' must be known at comptime.`);
         }
         if (typeof key.value !== 'string') {
@@ -1552,27 +1586,25 @@ ${stringifyNode(statement)}`);
         }
     }
 }
+const immutabilityReasons = {
+    constant: 'the left side is a constant',
+    'constant-immutable-def': 'the left side is a constant',
+    'runtime-immutable-def': 'the left side is immutable',
+    uniform: 'uniform buffers cannot be mutated',
+    readonly: 'readonly buffers cannot be mutated',
+    immediate: 'immediate variables cannot be mutated',
+    argument: 'non-pointer arguments cannot be mutated',
+    handle: 'textures and samplers cannot be mutated',
+};
 function validateSnippetMutation(mutated, expr) {
-    if (mutated.origin === 'constant' ||
-        mutated.origin === 'constant-immutable-def' ||
-        mutated.origin === 'runtime-immutable-def') {
-        if (isKnownAtComptime(mutated)) {
-            throw new WgslTypeError(`'${stringifyNode(expr)}' is invalid, because the left side is defined outside of the shader, and therefore is immutable during its execution. Try using tgpu.privateVar or buffers.`);
-        }
-        throw new WgslTypeError(`'${stringifyNode(expr)}' is invalid, because the left side is a constant.`);
-    }
-    if (mutated.origin === 'uniform') {
-        throw new WgslTypeError(`'${stringifyNode(expr)}' is invalid, because uniform buffers cannot be mutated.`);
-    }
-    if (mutated.origin === 'readonly') {
-        throw new WgslTypeError(`'${stringifyNode(expr)}' is invalid, because readonly buffers cannot be mutated.`);
-    }
-    if (mutated.origin === 'immediate') {
-        throw new WgslTypeError(`'${stringifyNode(expr)}' is invalid, because immediate variables cannot be mutated.`);
-    }
-    if (mutated.origin === 'argument') {
-        throw new WgslTypeError(`'${stringifyNode(expr)}' is invalid, because non-pointer arguments cannot be mutated.`);
+    if (isMutableOrigin(mutated.origin)) {
+        return;
     }
+    const reason = asComptime(mutated)
+        ? 'the left side is immutable during shader execution. Try using tgpu.privateVar or buffers'
+        : (immutabilityReasons[mutated.origin] ??
+            'the left side is not a reference to an existing value');
+    throw new WgslTypeError(`'${stringifyNode(expr)}' is invalid, because ${reason}.`);
 }
 function assertExhaustive(value) {
     throw new Error(`'${safeStringify(value)}' was not handled by the WGSL generator.`);

package/types.d.ts

@@ -18,7 +18,7 @@ import type { TgpuVar } from './core/variable/tgpuVariable.ts';
 import { type AnyData, UnknownData } from './data/dataTypes.ts';
 import type { MapValueToSnippet, ResolvedSnippet, Snippet } from './data/snippet.ts';
 import { type AnyMatInstance, type AnyVecInstance, type BaseData } from './data/wgslTypes.ts';
-import { $cast, $gpuCallable, $gpuValueOf, $internal, $ownSnippet, $resolve } from './shared/symbols.ts';
+import { $cast, $comptimeValueOf, $gpuCallable, $gpuValueOf, $internal, $ownSnippet, $resolve } from './shared/symbols.ts';
 import type { TgpuBindGroupLayout, TgpuLayoutEntry } from './tgpuBindGroupLayout.ts';
 import type { WgslEnableExtension } from './wgslExtensions.ts';
 import type { Infer } from './shared/repr.ts';
@@ -342,6 +342,23 @@ export interface WithOwnSnippet {
     readonly [$ownSnippet]: Snippet;
 }
 export declare function getOwnSnippet(value: unknown): Snippet | undefined;
+export type ComptimeValue = {
+    comptime: false;
+} | {
+    comptime: true;
+    value: unknown;
+};
+export interface WithComptimeValue {
+    /**
+     * @returns The value that this object represents at comptime,
+     * or `{ comptime: false }` if it is only known at runtime.
+     */
+    [$comptimeValueOf](): ComptimeValue;
+}
+/**
+ * A `[$comptimeValueOf]` implementation for resources that are only known at runtime
+ */
+export declare function runtimeOnly(): ComptimeValue;
 export interface GPUCallable<TArgs extends unknown[] = unknown[]> {
     [$gpuCallable]: {
         strictSignature?: {
@@ -358,7 +375,11 @@ export type WithCast<T = BaseData> = GPUCallable<[v?: Infer<T>]> & {
 export declare function hasCast(value: unknown): value is WithCast;
 type AnyFn = (...args: never[]) => unknown;
 export type DualFn<T extends AnyFn> = T & GPUCallable<Parameters<T>>;
-export declare function isKnownAtComptime(snippet: Snippet): boolean;
+/**
+ * @returns A snippet holding the JS value that `snippet` represents at comptime,
+ * or `undefined` if it is only known at runtime.
+ */
+export declare function asComptime(snippet: Snippet): Snippet | undefined;
 export declare function isConstant(snippet: Snippet): boolean;
 export declare function isWgsl(value: unknown): value is Wgsl;
 export type BindableBufferUsage = 'uniform' | 'readonly' | 'mutable';

package/types.js

@@ -1,7 +1,8 @@
 import { isLazy, isProviding, isSlot, } from "./core/slot/slotTypes.js";
 import { UnknownData } from "./data/dataTypes.js";
+import { withValue } from "./data/snippet.js";
 import { isWgslData, } from "./data/wgslTypes.js";
-import { $cast, $gpuCallable, $gpuValueOf, $internal, $ownSnippet, $resolve, } from "./shared/symbols.js";
+import { $cast, $comptimeValueOf, $gpuCallable, $gpuValueOf, $internal, $ownSnippet, $resolve, } from "./shared/symbols.js";
 export class NormalState {
     type = 'normal';
 }
@@ -23,18 +24,42 @@ export function isSelfResolvable(value) {
 export function getOwnSnippet(value) {
     return value?.[$ownSnippet];
 }
+/**
+ * A `[$comptimeValueOf]` implementation for resources that are only known at runtime
+ */
+export function runtimeOnly() {
+    return { comptime: false };
+}
 export function isGPUCallable(value) {
     return !!value?.[$gpuCallable];
 }
 export function hasCast(value) {
     return !!value?.[$cast];
 }
-export function isKnownAtComptime(snippet) {
-    return ((typeof snippet.value !== 'string' || snippet.dataType === UnknownData) &&
-        getOwnSnippet(snippet.value) === undefined);
+/**
+ * @returns A snippet holding the JS value that `snippet` represents at comptime,
+ * or `undefined` if it is only known at runtime.
+ */
+export function asComptime(snippet) {
+    const { value } = snippet;
+    if (typeof value === 'string' && snippet.dataType !== UnknownData) {
+        // Resolved WGSL code
+        return undefined;
+    }
+    const getter = value?.[$comptimeValueOf];
+    if (!getter) {
+        // Objects that meet the criteria and didn't explicitly opt out of
+        // being known at comptime, are treated as known at comptime.
+        // That's mainly because there are far more objects that are eligible
+        // and we want to be treated as known at comptime than those that are
+        // special cases and need to opt out.
+        return snippet;
+    }
+    const result = getter.call(value);
+    return result.comptime ? withValue(result.value, snippet) : undefined;
 }
 export function isConstant(snippet) {
-    return (isKnownAtComptime(snippet) ||
+    return (asComptime(snippet) !== undefined ||
         snippet.origin === 'constant' ||
         snippet.origin === 'constant-immutable-def');
 }