typegpu 0.12.6 → 0.12.7

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, 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";
@@ -476,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-----`);
                 }
             }
@@ -707,23 +709,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));
@@ -779,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', 
@@ -958,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);
@@ -996,7 +1018,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;
@@ -1077,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>';
@@ -1278,7 +1300,8 @@ ${this.ctx.pre}else ${alternate}`,
                         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 +1366,48 @@ ${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;
+            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);
-            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 +1415,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 +1423,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 +1449,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 +1460,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,