typegpu 0.12.6 → 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, isAlias, 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) {
@@ -476,7 +502,9 @@ export class WgslGenerator {
                 validateSnippetMutation(convLhs, expression);
                 this.tryMarkModified(lhs);
                 // Compound assignment operators are okay, e.g. +=, -=, *=, /=, ...
-                if (op === '=' && isAlias(rhsExpr) && !wgsl.isNaturallyEphemeral(rhsExpr.dataType)) {
+                if (op === '=' &&
+                    isStoredInMemory(rhsExpr) &&
+                    !wgsl.isNaturallyEphemeral(rhsExpr.dataType)) {
                     throw new WgslTypeError(`'${stringifyNode(expression)}' is invalid, because references cannot be assigned.\n-----\nTry '${stringifyNode(lhs)} = ${this.ctx.resolve(unptr(rhsExpr.dataType)).value}(${stringifyNode(rhs)})' to copy the value instead.\n-----`);
                 }
             }
@@ -486,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}'.`
                         : ''}`);
                 }
@@ -660,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[]`
@@ -692,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)));
@@ -707,23 +735,43 @@ export class WgslGenerator {
             }
             if (wgsl.isWgslStruct(structType)) {
                 const entries = {};
+                const keysWithSideEffectsInSourceOrder = [];
                 for (const prop of properties) {
                     const key = resolveUniqueKey(prop);
                     const value = prop[1];
                     const propType = structType.propTypes[key];
                     if (propType === undefined) {
                         // Evaluate every field even if it gets stripped by the struct schema
-                        void this._expression(value);
+                        const expr = this._expression(value);
+                        if (expr.possibleSideEffects) {
+                            logger.warn('suspicious', `\
+Object property '${stringifyObjectProperty(prop)}' in '${stringifyNode(expression)}' does not exist on type '${String(structType)}'.
+The generated shader will omit it, so its runtime side effects will not occur.`);
+                        }
                         continue;
                     }
                     const expr = this._typedExpression(value, propType);
                     entries[key] = expr;
+                    if (expr.possibleSideEffects) {
+                        keysWithSideEffectsInSourceOrder.push(key);
+                    }
                 }
                 for (const key of Object.keys(structType.propTypes)) {
                     if (entries[key] === undefined) {
                         throw new WgslTypeError(`Missing property ${key} in object literal for struct ${structType}`);
                     }
                 }
+                const keysWithSideEffectsInSchemaOrder = Object.keys(structType.propTypes).filter((key) => entries[key].possibleSideEffects);
+                const changesSideEffectsOrder = keysWithSideEffectsInSourceOrder.some((key, index) => key !== keysWithSideEffectsInSchemaOrder[index]);
+                if (changesSideEffectsOrder) {
+                    logger.warn('suspicious', `\
+Properties with possible side effects in '${stringifyNode(expression)}' do not match '${String(structType)}' declaration order:
+
+  Source order:           [${keysWithSideEffectsInSourceOrder.join(', ')}]
+  Declaration order:      [${keysWithSideEffectsInSchemaOrder.join(', ')}]
+
+The generated shader will evaluate them in declaration order.`);
+                }
                 const convertedSnippets = convertStructValues(this.ctx, structType, entries);
                 return snip(stitch `${this.ctx.resolve(structType).value}(${convertedSnippets})`, structType, 
                 /* origin */ 'runtime', convertedSnippets.some((s) => s.possibleSideEffects));
@@ -767,8 +815,11 @@ export class WgslGenerator {
             // 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);
@@ -779,8 +830,8 @@ export class WgslGenerator {
                     !alt ||
                     consequent.possibleSideEffects ||
                     alternative.possibleSideEffects ||
-                    (isAlias(consequent) && !wgsl.isNaturallyEphemeral(consequent.dataType)) ||
-                    (isAlias(alternative) && !wgsl.isNaturallyEphemeral(alternative.dataType))) {
+                    (isStoredInMemory(consequent) && !wgsl.isNaturallyEphemeral(consequent.dataType)) ||
+                    (isStoredInMemory(alternative) && !wgsl.isNaturallyEphemeral(alternative.dataType))) {
                     throw new Error(`Ternary operator '${stringifyNode(expression)}' is invalid. For more complex branching, please use 'std.select' or if/else statements.`);
                 }
                 return snip(stitch `select(${alt}, ${con}, ${convertedTest})`, con.dataType, 'runtime', 
@@ -958,7 +1009,7 @@ export class WgslGenerator {
             // The existence of `expectedReturnType` implies a function shell, which in turn implies that the
             // value will be copied on return anyway
             !expectedReturnType &&
-                isAlias(returnSnippet) &&
+                isStoredInMemory(returnSnippet) &&
                 !wgsl.isNaturallyEphemeral(returnSnippet.dataType) &&
                 returnSnippet.origin !== 'local-def') {
                 const str = stringifyNode(returnNode);
@@ -996,7 +1047,7 @@ Try 'return ${typeStr}(${str});' instead.
             const declaration = `let ${rawId}`;
             throw new WgslTypeError(`'${declaration} = ${rhsStr}' is invalid, cannot determine WGSL type of '${rhsStr}'${schemaWrappingSuggestion(declaration, rhsStr, eq.value)}`);
         }
-        if (isAlias(eq) && !wgsl.isNaturallyEphemeral(eq.dataType)) {
+        if (isStoredInMemory(eq) && !wgsl.isNaturallyEphemeral(eq.dataType)) {
             // `let` declarations cannot store references
             const rhsStr = stringifyNode(eqNode);
             const rhsTypeStr = this.ctx.resolve(unptr(eq.dataType)).value;
@@ -1033,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 {
@@ -1077,7 +1130,7 @@ Try 'return ${typeStr}(${str});' instead.
             // function code has been processed, so at least currently, we lose that info.
             varOrigin = 'local-def';
         }
-        else if (!isAlias(eq)) {
+        else if (!isStoredInMemory(eq)) {
             // Not a reference, but also not naturally ephemeral, so we cannot guarantee it won't be mutated.
             // We defer the decision for now.
             varType = '<deferred>';
@@ -1265,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;
@@ -1273,12 +1327,13 @@ ${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 };
                     }
                     const firstElement = elements[0];
-                    if (!isAlias(firstElement) && !wgsl.isNaturallyEphemeral(firstElement.dataType)) {
+                    if (!isStoredInMemory(firstElement) &&
+                        !wgsl.isNaturallyEphemeral(firstElement.dataType)) {
                         throw new WgslTypeError(`Cannot unroll '${stringifyNode(iterable)}'. The elements of iterable are constructed in place but are not value types.`);
                     }
                     let blocksCode = '';
@@ -1343,11 +1398,50 @@ ${this.ctx.pre}else ${alternate}`,
         if (statement[0] === NODE.switch) {
             // Switch statement
             const [_, discriminant, cases] = statement;
-            const discriminantExpr = this._typedExpression(discriminant, [i32, u32]);
+            let discriminantExpr = this._expression(discriminant);
+            let caseExprs = cases.map(([test, consequent]) => {
+                const testExpr = test === null ? switchDefault : this._expression(test);
+                return [testExpr, consequent];
+            });
+            // comptime folding
+            let matchedCaseWasNotLast = false;
+            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);
+                }
+                if (matchedCaseIndex === -1) {
+                    return { code: '', definesInNearestScope: false };
+                }
+                const matchedConsequent = caseExprs
+                    .slice(matchedCaseIndex)
+                    .find(([, consequent]) => consequent.length !== 0)?.[1];
+                if (matchedConsequent === undefined) {
+                    return { code: '', definesInNearestScope: false };
+                }
+                if (matchedConsequent !== caseExprs.at(-1)?.[1]) {
+                    // The JS behavior will be different if this case does not end with control flow.
+                    matchedCaseWasNotLast = true;
+                }
+                discriminantExpr = snip(0, u32, 'constant', false);
+                caseExprs = [[switchDefault, matchedConsequent]];
+            }
+            // type concretization
+            discriminantExpr = tryConvertSnippet(this.ctx, discriminantExpr, [i32, u32]);
             const switchType = discriminantExpr.dataType;
             invariant(switchType !== UnknownData);
-            const caseExprs = cases.map(([test, consequent]) => {
-                const testExpr = test === null ? switchDefault : this._typedExpression(test, [switchType]);
+            caseExprs = caseExprs.map(([test, consequent]) => {
+                // TODO(#3148): Remove this throw.
+                if (test !== switchDefault && test.dataType === UnknownData) {
+                    throw new Error(`Failed to convert one of the switch tests to '${switchType.type}'`);
+                }
+                const convertedTest = test === switchDefault ? test : tryConvertSnippet(this.ctx, test, [switchType]);
+                return [convertedTest, consequent];
+            });
+            // consequent resolution
+            const resolvedCaseExprs = caseExprs.map(([test, consequent]) => {
                 // In WGSL, each case is a different block. This block scope forbids scope leaking.
                 // TODO(#3001): Consider using NODE.block here
                 this.ctx.pushBlockScope();
@@ -1355,7 +1449,7 @@ ${this.ctx.pre}else ${alternate}`,
                 this.ctx.indent();
                 try {
                     const consequentStmts = consequent.map((s) => this._statement(s));
-                    return [testExpr, consequentStmts];
+                    return [test, consequentStmts];
                 }
                 finally {
                     this.ctx.dedent();
@@ -1363,6 +1457,15 @@ ${this.ctx.pre}else ${alternate}`,
                     this.ctx.popBlockScope();
                 }
             });
+            const groupedCaseExprs = [];
+            let currentGroup = [];
+            for (const [index, [test, consequent]] of resolvedCaseExprs.entries()) {
+                currentGroup.push(test);
+                if (consequent.length > 0 || index === caseExprs.length - 1) {
+                    groupedCaseExprs.push([currentGroup, consequent]);
+                    currentGroup = [];
+                }
+            }
             // Validation
             {
                 // Tests should be constant
@@ -1380,7 +1483,9 @@ ${stringifyNode(statement)}`);
                 // We skip this check, because WGSL errors are readable,
                 // and we cannot easily access non-comptime known constants.
                 // Tests should not have non-trivial fallthrough
-                caseExprs.slice(0, -1).forEach(([_, consequent]) => {
+                resolvedCaseExprs
+                    .slice(0, matchedCaseWasNotLast ? undefined : -1)
+                    .forEach(([_, consequent]) => {
                     const last = consequent.at(-1);
                     if (last && !last.endsWithControlFlow) {
                         throw new Error(`Switch statement cannot have non-trivial fallthrough.
@@ -1389,15 +1494,6 @@ ${stringifyNode(statement)}`);
                     }
                 });
             }
-            const groupedCaseExprs = [];
-            let currentGroup = [];
-            for (const [index, [test, consequent]] of caseExprs.entries()) {
-                currentGroup.push(test);
-                if (consequent.length > 0 || index === caseExprs.length - 1) {
-                    groupedCaseExprs.push([currentGroup, consequent]);
-                    currentGroup = [];
-                }
-            }
             return {
                 code: this._emitSwitchStatement(discriminantExpr, groupedCaseExprs),
                 definesInNearestScope: false,
@@ -1455,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') {
@@ -1490,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');
 }