typegpu 0.12.5 → 0.12.7

package/tgsl/wgslGenerator.js

@@ -4,16 +4,16 @@ 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, isStoredInMemory, snip, } from "../data/snippet.js";
 import * as wgsl from "../data/wgslTypes.js";
-import { invariant, ResolutionError, WgslTypeError } from "../errors.js";
+import { invariant, ResolutionError, WgslForbiddenStatementError, WgslTypeError, } from "../errors.js";
 import { getName } from "../shared/meta.js";
 import { $gpuCallable, $internal, $providing, isMarkedInternal } from "../shared/symbols.js";
 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 { isGPUCallable, isKnownAtComptime, } from "../types.js";
+import { isConstant, isGPUCallable, isKnownAtComptime, } from "../types.js";
 import { convertStructValues, convertToCommonType, tryConvertSnippet } from "./conversion.js";
 import { ArrayExpression, coerceToSnippet, concretize, numericLiteralToSnippet, } from "./generationHelpers.js";
 import { accessIndex } from "./accessIndex.js";
@@ -28,11 +28,12 @@ import { AutoStruct } from "../data/autoStruct.js";
 import { mathToStd, supportedLogOps } from "./jsPolyfills.js";
 import * as forOfUtils from "./forOfUtils.js";
 import { isTgpuRange } from "../std/range.js";
-import { stringifyNode } from "../shared/tseynit.js";
+import { stringifyNode, stringifyObjectProperty } from "../shared/tseynit.js";
 import { getAttributesString } from "../data/attributes.js";
 import { validSelectBranchTypes } from "../std/boolean.js";
 import { isInfixDispatch } from "./infixDispatch.js";
 import { logger } from "../tgpuLogger.js";
+import { TgpuDeclareImpl } from "../core/declare/tgpuDeclare.js";
 const { NodeTypeCatalog: NODE } = tinyest;
 const parenthesizedOps = [
     '==',
@@ -151,11 +152,23 @@ const usageToVarTemplateMap = {
     uniform: 'uniform',
     mutable: 'storage, read_write',
     readonly: 'storage, read',
+    immediate: 'immediate',
 };
 /**
  * The block depth that we can expect when generating code in the function scope, not in any nested blocks.
  */
 const functionInitialBlockDepth = 2;
+// Used for switch case.
+const switchDefault = snip('default', UnknownData, 'constant');
+function schemaWrappingSuggestion(declaration, rhs, value) {
+    if (value === null || value === undefined || typeof value === 'string') {
+        return '';
+    }
+    return `
+-----
+- Try using or defining a schema that matches your desired value the most, and wrap the value with it: '${declaration} = Schema(${rhs})'
+-----`;
+}
 export class WgslGenerator {
     #ctx = undefined;
     // used to detect `continue` and `break` nodes in loop body, as well as label
@@ -265,6 +278,27 @@ 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).
+            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}}`;
+        });
+        this.ctx.dedent();
+        return stitch `${this.ctx.pre}switch ${discriminantExpr} {\n${cases.join('\n')}\n${this.ctx.pre}}`;
+    }
     _callShellless(callee, args) {
         const isGeneric = isGenericFn(callee);
         const slotPairs = isGeneric ? (callee[$providing]?.pairs ?? []) : [];
@@ -308,11 +342,11 @@ export class WgslGenerator {
         }
     }
     _expression(expression) {
-        if (typeof expression === 'string') {
-            return this._identifier(expression);
+        if (isId(expression)) {
+            return this._identifier(extractId(expression));
         }
-        if (typeof expression === 'boolean') {
-            return snip(expression, bool, /* origin */ 'constant', false);
+        if (isBool(expression)) {
+            return snip(extractBool(expression), bool, /* origin */ 'constant', false);
         }
         if (expression[0] === NODE.logicalExpr) {
             const [_, lhs, op, rhs] = expression;
@@ -442,7 +476,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-----`);
                 }
             }
@@ -481,8 +517,9 @@ export class WgslGenerator {
         }
         if (expression[0] === NODE.memberAccess) {
             // Member Access
-            const [_, targetNode, property] = expression;
+            const [_, targetNode, propertyNode] = expression;
             const target = this._expression(targetNode);
+            const property = extractId(propertyNode);
             const accessed = accessProp(target, property);
             if (!accessed) {
                 throw new Error(`Property '${property}' not found on '${stringifyNode(targetNode)}'`);
@@ -560,7 +597,7 @@ export class WgslGenerator {
                 // We just let the argument resolve everything.
                 return snip(this.ctx.resolveSnippet(arg).value, callee.value, 
                 // A new array, so not a reference.
-                /* origin */ 'runtime', arg.possibleSideEffects);
+                /* origin */ fallthroughCopyOrigin(arg.origin), arg.possibleSideEffects);
             }
             if (callee.value === constant) {
                 throw new Error('Constants cannot be defined within TypeGPU function scope. To address this, move the constant definition outside the function scope.');
@@ -628,11 +665,25 @@ export class WgslGenerator {
             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`);
         }
         if (expression[0] === NODE.objectExpr) {
-            // Object Literal
-            const obj = expression[1];
+            // Normalize to `objectProperty[]`
+            const properties = Array.isArray(expression[1])
+                ? expression[1]
+                : Object.entries(expression[1]).map(([key, value]) => [key, value, false]);
+            const seenKeys = new Map();
+            const resolveUniqueKey = (prop) => {
+                const key = this._resolveObjectPropertyKey(prop);
+                const dupProp = seenKeys.get(key);
+                if (dupProp) {
+                    throw new WgslTypeError(`Duplicate object property key found: '${stringifyObjectProperty(dupProp)}' and '${stringifyObjectProperty(prop)}'.`);
+                }
+                seenKeys.set(key, prop);
+                return key;
+            };
             const structType = this.ctx.expectedType;
             if (structType instanceof AutoStruct) {
-                const entries = Object.fromEntries(Object.entries(obj).map(([key, value]) => {
+                const keySnippetPairs = properties.map((prop) => {
+                    const key = resolveUniqueKey(prop);
+                    const value = prop[1];
                     let accessed = structType.accessProp(key);
                     let expr;
                     if (accessed) {
@@ -649,21 +700,52 @@ export class WgslGenerator {
                         accessed = structType.provideProp(key, unptr(concretize(expr.dataType)));
                     }
                     return [accessed.prop, expr];
-                }));
+                });
+                const entries = Object.fromEntries(keySnippetPairs);
                 const completeStruct = structType.completeStruct;
                 const convertedSnippets = convertStructValues(this.ctx, completeStruct, entries);
                 return snip(stitch `${this.ctx.resolve(structType).value}(${convertedSnippets})`, completeStruct, 
                 /* origin */ 'runtime');
             }
             if (wgsl.isWgslStruct(structType)) {
-                const entries = Object.fromEntries(Object.entries(structType.propTypes).map(([key, value]) => {
-                    const val = obj[key];
-                    if (val === undefined) {
+                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
+                        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 result = this._typedExpression(val, value);
-                    return [key, result];
-                }));
+                }
+                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));
@@ -719,8 +801,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', 
@@ -818,8 +900,7 @@ export class WgslGenerator {
             // We can make this snippet ephemeral, as we know it will be deep copied in JS
             return snip(stitch `${args[0]}`, schema, fallthroughCopyOrigin(args[0].origin), args[0].possibleSideEffects);
         }
-        // Creating a 'runtime' snippet, since it's instantiating a new value
-        return snip(stitch `${this.ctx.resolve(schema).value}(${args})`, schema, 'runtime', args.some((s) => s.possibleSideEffects));
+        return snip(stitch `${this.ctx.resolve(schema).value}(${args})`, schema, args.every((arg) => arg.origin === 'constant') ? 'constant' : 'runtime', args.some((s) => s.possibleSideEffects));
     }
     numericLiteral(value, schema) {
         if (!Number.isFinite(value)) {
@@ -899,7 +980,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);
@@ -918,7 +999,8 @@ Try 'return ${typeStr}(${str});' instead.
         return `${this.ctx.pre}return;`;
     }
     _letStatement(statement) {
-        const [_, rawId, eqNode] = statement;
+        const [_, rawIdNode, eqNode] = statement;
+        const rawId = extractId(rawIdNode);
         if (eqNode === undefined) {
             throw new Error(`'${stringifyNode(statement)}' is invalid because all variables need initializers.`);
         }
@@ -931,14 +1013,12 @@ Try 'return ${typeStr}(${str});' instead.
 -----`);
         }
         const definitionDataType = eq.dataType;
-        if (definitionDataType === UnknownData) {
+        if (definitionDataType === UnknownData || wgsl.isVoid(definitionDataType)) {
             const rhsStr = stringifyNode(eqNode);
-            throw new WgslTypeError(`'let ${rawId} = ${rhsStr}' is invalid, cannot determine WGSL type of '${rhsStr}'
------
-- Try using or defining a schema that matches your desired value the most, and wrap the value with it: 'let ${rawId} = Schema(${rhsStr})'
------`);
+            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;
@@ -967,7 +1047,8 @@ Try 'return ${typeStr}(${str});' instead.
         };
     }
     _constStatement(statement) {
-        const [_, rawId, eqNode] = statement;
+        const [_, rawIdNode, eqNode] = statement;
+        const rawId = extractId(rawIdNode);
         if (eqNode === undefined) {
             throw new Error(`'${stringifyNode(statement)}' is invalid because all variables need initializers.`);
         }
@@ -988,12 +1069,10 @@ Try 'return ${typeStr}(${str});' instead.
         let varOrigin = 'local-def';
         let varType = '<deferred>';
         let definitionDataType = eq.dataType;
-        if (definitionDataType === UnknownData) {
+        if (definitionDataType === UnknownData || wgsl.isVoid(definitionDataType)) {
             const rhsStr = stringifyNode(eqNode);
-            throw new WgslTypeError(`'const ${rawId} = ${rhsStr}' is invalid, cannot determine WGSL type of '${rhsStr}'
------
-- Try using or defining a schema that matches your desired value the most, and wrap the value with it: 'const ${rawId} = Schema(${rhsStr})'
------`);
+            const declaration = `const ${rawId}`;
+            throw new WgslTypeError(`'${declaration} = ${rhsStr}' is invalid, cannot determine WGSL type of '${rhsStr}'${schemaWrappingSuggestion(declaration, rhsStr, eq.value)}`);
         }
         if (eq.origin === 'argument') {
             // Arguments are immutable, so we 'let' them be (kill me)
@@ -1020,7 +1099,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>';
@@ -1083,16 +1162,16 @@ Try 'return ${typeStr}(${str});' instead.
         };
     }
     _statement(statement) {
-        if (typeof statement === 'string') {
-            const id = this._identifier(statement);
-            const resolved = id.value !== undefined && id.value !== null ? this.ctx.resolveSnippet(id).value : '';
-            return { code: resolved ? `${this.ctx.pre}${resolved};` : '', definesInNearestScope: false };
+        // TODO(#3078): Remove this check.
+        if (isId(statement)) {
+            const item = this.ctx.getById(extractId(statement));
+            if (item?.value instanceof TgpuDeclareImpl) {
+                this.ctx.resolveSnippet(item);
+                return { code: '', definesInNearestScope: false };
+            }
         }
-        if (typeof statement === 'boolean') {
-            return {
-                code: `${this.ctx.pre}${statement ? 'true' : 'false'};`,
-                definesInNearestScope: false,
-            };
+        if (isId(statement) || isBool(statement)) {
+            throw new WgslForbiddenStatementError(statement);
         }
         if (statement[0] === NODE.return) {
             return {
@@ -1205,24 +1284,24 @@ ${this.ctx.pre}else ${alternate}`,
                 const shouldUnroll = iterableExpr.value instanceof UnrollableIterable;
                 const iterableSnippet = shouldUnroll ? iterableExpr.value.snippet : iterableExpr;
                 const range = forOfUtils.getRangeSnippets(this.ctx, iterableSnippet, shouldUnroll);
-                const originalLoopVarName = loopVar[1];
+                const originalLoopVarName = extractId(loopVar[1]);
                 const blockified = blockifySingleStatement(body);
                 if (shouldUnroll) {
                     if (!isKnownAtComptime(range.end)) {
                         throw new Error('Cannot unroll loop. Length of iterable is unknown at comptime.');
                     }
-                    const length = range.end.value;
-                    if (length === 0) {
-                        return { code: '', definesInNearestScope: false };
-                    }
                     const { value } = iterableSnippet;
                     const elements = isTgpuRange(value)
                         ? value.map((i) => coerceToSnippet(i))
                         : value instanceof ArrayExpression
                             ? value.elements
-                            : Array.from({ length }, (_, i) => forOfUtils.getElementSnippet(iterableSnippet, snip(i, u32, 'constant')));
+                            : Array.from({ length: range.end.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 = '';
@@ -1284,6 +1363,108 @@ ${this.ctx.pre}else ${alternate}`,
                 this.ctx.popBlockScope();
             }
         }
+        if (statement[0] === NODE.switch) {
+            // Switch statement
+            const [_, discriminant, cases] = statement;
+            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;
+            if ([discriminantExpr, ...caseExprs.map(([test]) => test)].every(isKnownAtComptime)) {
+                let matchedCaseIndex = caseExprs.findIndex(([test]) => test !== switchDefault && test.value === discriminantExpr.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);
+            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();
+                this.ctx.indent();
+                this.ctx.indent();
+                try {
+                    const consequentStmts = consequent.map((s) => this._statement(s));
+                    return [test, consequentStmts];
+                }
+                finally {
+                    this.ctx.dedent();
+                    this.ctx.dedent();
+                    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
+                caseExprs.forEach(([testExpr], i) => {
+                    if (!isConstant(testExpr)) {
+                        const testNode = cases[i]?.[0];
+                        invariant(testNode, `Expected node to be not nullish.`);
+                        throw new Error(`All of switch tests must be constant.
+Test '${stringifyNode(testNode)}' is not constant, making the following switch statement invalid. 
+This error may be caused by an implicit conversion.
+${stringifyNode(statement)}`);
+                    }
+                });
+                // Tests should not have duplicates.
+                // 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
+                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.
+The following switch statement is invalid:
+${stringifyNode(statement)}`);
+                    }
+                });
+            }
+            return {
+                code: this._emitSwitchStatement(discriminantExpr, groupedCaseExprs),
+                definesInNearestScope: false,
+            };
+        }
         if (statement[0] === NODE.postUpdate) {
             // Post-update statement
             const [_, op, arg] = statement;
@@ -1313,10 +1494,38 @@ ${this.ctx.pre}else ${alternate}`,
                 definesInNearestScope: false,
             };
         }
+        if (statement[0] === NODE.numericLiteral ||
+            statement[0] === NODE.binaryExpr ||
+            statement[0] === NODE.unaryExpr ||
+            statement[0] === NODE.logicalExpr ||
+            statement[0] === NODE.arrayExpr ||
+            statement[0] === NODE.objectExpr ||
+            statement[0] === NODE.nullLiteral) {
+            // `memberAccess` is allowed for raw code snippets, `indexAccess` is allowed for consistency.
+            // `stringLiteral` is forbidden by our injection prevention anyway.
+            throw new WgslForbiddenStatementError(statement);
+        }
         const expr = this._expression(statement);
         const resolved = expr.value !== undefined && expr.value !== null ? this.ctx.resolveSnippet(expr).value : '';
         return { code: resolved ? `${this.ctx.pre}${resolved};` : '', definesInNearestScope: false };
     }
+    /**
+     * Resolves the key of an object property. Handles both computed and non-computed keys.
+     */
+    _resolveObjectPropertyKey(property) {
+        const computed = property[2];
+        if (!computed) {
+            return property[0];
+        }
+        const key = this._expression(property[0]);
+        if (!isKnownAtComptime(key)) {
+            throw new WgslTypeError(`Computed object property key '${stringifyObjectProperty(property)}' must be known at comptime.`);
+        }
+        if (typeof key.value !== 'string') {
+            throw new WgslTypeError('Object property keys must be strings in TypeGPU functions.');
+        }
+        return key.value;
+    }
     /**
      * Attempts a member access lookup to mark a variable as modified.
      * @example
@@ -1358,6 +1567,9 @@ function validateSnippetMutation(mutated, expr) {
     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.`);
     }
@@ -1387,7 +1599,25 @@ function extractObject(expr) {
         (object[0] === NODE.memberAccess || object[0] === NODE.indexAccess)) {
         object = object[1];
     }
-    if (typeof object === 'string') {
-        return object;
+    if (isId(object)) {
+        return extractId(object);
+    }
+}
+function isId(expr) {
+    return typeof expr === 'string' || (Array.isArray(expr) && expr[0] === NODE.identifier);
+}
+function extractId(ident) {
+    if (typeof ident === 'string') {
+        return ident;
+    }
+    return ident[1];
+}
+function isBool(expr) {
+    return typeof expr === 'boolean' || (Array.isArray(expr) && expr[0] === NODE.booleanLiteral);
+}
+function extractBool(ident) {
+    if (typeof ident === 'boolean') {
+        return ident;
     }
+    return ident[1];
 }

package/types.d.ts

@@ -13,6 +13,7 @@ import type { TgpuSampler } from './core/sampler/sampler.ts';
 import { type Eventual, type SlotValuePair, type TgpuAccessor, type TgpuLazy, type TgpuSlot } from './core/slot/slotTypes.ts';
 import type { TgpuExternalTexture } from './core/texture/externalTexture.ts';
 import type { TgpuTexture, TgpuTextureView } from './core/texture/texture.ts';
+import type { TgpuImmediateVar } from './core/immediate/immediateVar.ts';
 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';
@@ -26,7 +27,7 @@ import type { StorageFlag } from './extension.ts';
 import type { TgpuBufferBinding } from './core/buffer/bufferBinding.ts';
 import type { ShelllessRepository } from './tgsl/shellless.ts';
 import type { SupportedLogOp } from './tgsl/consoleLog/types.ts';
-export type ResolvableObject = SelfResolvable | TgpuLazy<unknown> | TgpuConst | TgpuDeclare | TgpuBindGroupLayout | TgpuFn | TgpuComputeFn | TgpuFragmentFn | TgpuComputePipeline | TgpuRenderPipeline | TgpuVertexFn | TgpuSampler | TgpuAccessor | TgpuExternalTexture | TgpuTexture | TgpuTextureView | TgpuBufferBinding<BaseData> | TgpuVar | AnyVecInstance | AnyMatInstance | AnyData | ((...args: never[]) => unknown);
+export type ResolvableObject = SelfResolvable | TgpuLazy<unknown> | TgpuConst | TgpuDeclare | TgpuBindGroupLayout | TgpuFn | TgpuComputeFn | TgpuFragmentFn | TgpuComputePipeline | TgpuRenderPipeline | TgpuVertexFn | TgpuSampler | TgpuAccessor | TgpuExternalTexture | TgpuTexture | TgpuTextureView | TgpuBufferBinding<BaseData> | TgpuVar | TgpuImmediateVar | AnyVecInstance | AnyMatInstance | AnyData | ((...args: never[]) => unknown);
 export type Wgsl = Eventual<number | boolean | ResolvableObject>;
 export type ShaderStage = 'compute' | 'vertex' | 'fragment';
 export interface ResolveFunctionOptions {
@@ -281,6 +282,12 @@ export interface ResolutionCtx {
      * @returns an identifier that is unique within the given scope
      */
     makeUniqueIdentifier(primer: string | undefined, scope: 'global' | 'block'): string;
+    /**
+     * Registers the use of an immediate variable in the current resolution.
+     * At most one `var<immediate>` is allowed per resolution.
+     * @throws When a different immediate variable has already been registered.
+     */
+    registerImmediate(immediate: TgpuImmediateVar): void;
     isIdentifierBanned(name: string): boolean;
     /**
      * @param name The name to check.
@@ -352,6 +359,7 @@ 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;
+export declare function isConstant(snippet: Snippet): boolean;
 export declare function isWgsl(value: unknown): value is Wgsl;
 export type BindableBufferUsage = 'uniform' | 'readonly' | 'mutable';
 export type BufferUsage = 'uniform' | 'readonly' | 'mutable' | 'vertex';

package/types.js

@@ -33,6 +33,11 @@ export function isKnownAtComptime(snippet) {
     return ((typeof snippet.value !== 'string' || snippet.dataType === UnknownData) &&
         getOwnSnippet(snippet.value) === undefined);
 }
+export function isConstant(snippet) {
+    return (isKnownAtComptime(snippet) ||
+        snippet.origin === 'constant' ||
+        snippet.origin === 'constant-immutable-def');
+}
 export function isWgsl(value) {
     return (typeof value === 'number' ||
         typeof value === 'boolean' ||