@woosh/meep-engine 3.34.0 → 3.35.0

package/src/shade/wgsl/emulator/WGSLJavaScriptCompiler.d.ts

@@ -32,7 +32,7 @@ export class WGSLJavaScriptCompiler {
     compile(ast: Tree): string;
 }
 export type WgslType = {
-    kind: 'primitive' | 'vec' | 'mat' | 'array' | 'struct' | 'texture_2d' | 'sampler' | 'atomic' | 'unknown';
+    kind: 'primitive' | 'abstract' | 'vec' | 'mat' | 'array' | 'struct' | 'texture_2d' | 'sampler' | 'atomic' | 'unknown';
     name?: string;
     size?: number;
     cols?: number;

package/src/shade/wgsl/emulator/WGSLJavaScriptCompiler.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"WGSLJavaScriptCompiler.d.ts","sourceRoot":"","sources":["../../../../../src/shade/wgsl/emulator/WGSLJavaScriptCompiler.js"],"names":[],"mappings":"AA0hBA;IAEI;;;;;;;;;;;;;;;;;;;;;;;;;;;;;OA6BG;IACH,oBAFa,MAAM,CAMlB;CACJ;;UAtjBa,WAAW,GAAG,KAAK,GAAG,KAAK,GAAG,OAAO,GAAG,QAAQ,GAAG,YAAY,GAAG,SAAS,GAAG,QAAQ,GAAG,SAAS;WAClG,MAAM;WACN,MAAM;WACN,MAAM;WACN,MAAM;aACN,MAAM;cACN,QAAQ;YACR,MAAM"}
+{"version":3,"file":"WGSLJavaScriptCompiler.d.ts","sourceRoot":"","sources":["../../../../../src/shade/wgsl/emulator/WGSLJavaScriptCompiler.js"],"names":[],"mappings":"AAuiBA;IAEI;;;;;;;;;;;;;;;;;;;;;;;;;;;;;OA6BG;IACH,oBAFa,MAAM,CAMlB;CACJ;;UAnkBa,WAAW,GAAG,UAAU,GAAG,KAAK,GAAG,KAAK,GAAG,OAAO,GAAG,QAAQ,GAAG,YAAY,GAAG,SAAS,GAAG,QAAQ,GAAG,SAAS;WAC/G,MAAM;WACN,MAAM;WACN,MAAM;WACN,MAAM;aACN,MAAM;cACN,QAAQ;YACR,MAAM"}

package/src/shade/wgsl/emulator/WGSLJavaScriptCompiler.js

@@ -7,7 +7,7 @@ import LineBuilder from "../../../core/codegen/LineBuilder.js";
 
 /**
  * @typedef {object} WgslType
- * @property {'primitive' | 'vec' | 'mat' | 'array' | 'struct' | 'texture_2d' | 'sampler' | 'atomic' | 'unknown'} kind
+ * @property {'primitive' | 'abstract' | 'vec' | 'mat' | 'array' | 'struct' | 'texture_2d' | 'sampler' | 'atomic' | 'unknown'} kind
  * @property {string} [name]
  * @property {number} [size]
  * @property {number} [cols]
@@ -19,6 +19,19 @@ import LineBuilder from "../../../core/codegen/LineBuilder.js";
 
 const T_UNKNOWN = Object.freeze({ kind: 'unknown' });
 const T_I32 = Object.freeze({ kind: 'primitive', name: 'i32' });
+const T_F32 = Object.freeze({ kind: 'primitive', name: 'f32' });
+
+/**
+ * WGSL's two abstract numeric types. A bare literal has one (`2` is an AbstractInt, `0.5` an
+ * AbstractFloat), arithmetic between abstract operands keeps one, and so does a `const` declared
+ * without a type — at any scope. An abstract value becomes concrete where it meets a concrete
+ * operand, which gives it that operand's type, or where it is bound, passed, returned or stored.
+ * See {@link compileExpression} and {@link concretizeOperands}.
+ *
+ * An AbstractFloat is a binary64, so a const-expression of them is plain JS arithmetic.
+ */
+const T_ABSTRACT_INT = Object.freeze({ kind: 'abstract', name: 'AbstractInt' });
+const T_ABSTRACT_FLOAT = Object.freeze({ kind: 'abstract', name: 'AbstractFloat' });
 
 const VECTOR_ALIASES = {
     vec2i: { size: 2, scalar: 'i32' },
@@ -328,9 +341,9 @@ function literalToJs(text) {
 /**
  * The tracked type of a literal token, as far as its own spelling settles it.
  *
- * Only an explicit suffix is honoured. A bare `2043` is WGSL's AbstractInt, which takes its
- * concrete type from context this compiler does not track, so it stays unknown rather than being
- * guessed at.
+ * A suffix names a concrete type. Without one the literal is abstract: a bare `2043` is an
+ * AbstractInt and a bare `0.5` an AbstractFloat, and each takes a concrete type from where it is
+ * used (see {@link T_ABSTRACT_INT}).
  *
  * The suffix matters because operator selection reads these types. `index / 2043u` is truncating
  * integer division in WGSL and plain division in JavaScript, and the two differ for every pair
@@ -359,10 +372,10 @@ function literalType(text) {
     }
 
     if (/[.eE]/.test(t) && !/^0[xX]/.test(t)) {
-        return { kind: 'primitive', name: 'f32' };
+        return T_ABSTRACT_FLOAT;
     }
 
-    return T_UNKNOWN;
+    return T_ABSTRACT_INT;
 }
 
 // =========================================================================
@@ -1025,13 +1038,20 @@ function collectGlobalValue(c, node) {
     const name = ident.text;
     let type = type_node !== null ? parseTypeSpecifier(type_node) : T_UNKNOWN;
 
-    const expr = compileExpression(c, value_node);
+    // A `const` without a type stays abstract, as it does in a function. An `override` is always
+    // concrete, and an untyped one concretizes as a `let` does.
+    const is_override = /\boverride\b/.test(node.text.slice(0, ident.startIndex - node.startIndex));
+    const keeps_abstract = type_node === null && !is_override;
+
+    const expr = keeps_abstract
+        ? compileExpressionAbstract(c, value_node)
+        : compileExpression(c, value_node);
 
     // Infer the type from the RHS when there's no explicit annotation (mirrors the local
     // `let`/`var` path). Without this a `const M = mat3x3<f32>(...)` is left T_UNKNOWN, so a
     // later `M * v` can't tell M is a matrix and emits a vector multiply → length mismatch.
     if (type === T_UNKNOWN || type.kind === 'unknown') {
-        type = expr.type;
+        type = keeps_abstract ? expr.type : concreteDefault(expr.type);
     }
 
     c.scope.declare(name, type);
@@ -1263,6 +1283,7 @@ function emitVariableOrValueStatement(c, node) {
     let kind = 'let';
     let opt_ident = null;
     let template = null;
+    let is_wgsl_const = false;
 
     if (var_decl !== null) {
         // WGSL `var` is the mutable form → JS `let`.
@@ -1274,6 +1295,7 @@ function emitVariableOrValueStatement(c, node) {
         // Tree-sitter exposes the keyword as an unnamed leading token; read it from the source.
         kind = 'const';
         opt_ident = opt_ident_direct;
+        is_wgsl_const = textBeforeFirstNamed(node).trim() === 'const';
     }
 
     const ident = namedChildByType(opt_ident, 'ident');
@@ -1282,11 +1304,15 @@ function emitVariableOrValueStatement(c, node) {
     let type = type_spec !== null ? parseTypeSpecifier(type_spec) : T_UNKNOWN;
 
     if (expr_node !== null) {
-        const expr = compileExpression(c, expr_node);
-        // Infer the type from the RHS when no explicit annotation is present. A bare integer
-        // literal concretises to i32 here, as WGSL does for `let x = 3;`.
+        // A `const` without a type keeps an abstract initializer abstract: `const A = 2;` is an
+        // AbstractInt, and `A * x` with an f32 x is an f32 product. A `let` or `var` concretizes
+        // it, so `let x = 3;` is an i32 and `let h = 0.5;` an f32.
+        const keeps_abstract = is_wgsl_const && type_spec === null;
+        const expr = keeps_abstract
+            ? compileExpressionAbstract(c, expr_node)
+            : compileExpression(c, expr_node);
         if (type === T_UNKNOWN || type.kind === 'unknown') {
-            type = expr.abstract_int === true ? T_I32 : expr.type;
+            type = keeps_abstract ? expr.type : concreteDefault(expr.type);
         }
         c.scope.declare(name, type);
         c.out.add(`${kind} ${identToJs(name)} = ${expr.code};`);
@@ -1542,7 +1568,7 @@ function emitForInitInline(c, node) {
         let type = type_spec !== null ? parseTypeSpecifier(type_spec) : T_UNKNOWN;
         if (type.kind === 'unknown' && init !== null) {
             // `for (var i = 0; ...)`: a bare integer literal concretises to i32, as in a `let`
-            type = init.abstract_int === true ? T_I32 : init.type;
+            type = concreteDefault(init.type);
         }
         c.scope.declare(name, type);
         const value = init !== null ? init.code : defaultValueForType(type);
@@ -1758,12 +1784,38 @@ function emitFunctionCallStatement(c, node) {
 // =========================================================================
 
 /**
- * Compile an expression node (any wrapper or specific kind).
+ * Compile an expression whose value is used as a concrete one: bound, passed, returned, stored,
+ * indexed with, branched on.
+ *
+ * An AbstractFloat result becomes an f32 here, which is what a WGSL AbstractFloat converts to
+ * everywhere but an f16 context. Doing it by default, rather than at each place WGSL converts,
+ * means a use this compiler has not thought of rounds like the device rather than staying
+ * more accurate than it. An AbstractInt result is left as it is: its value is the same in every
+ * type it can take, and a `let` still has to see it is abstract to make it an i32.
+ *
+ * {@link compileExpressionAbstract} is the form for the few places an abstract value stays
+ * abstract: the operands of an operator, and the initializer of a `const` without a type.
+ *
  * @param {Compiler} c
  * @param {Node} node
- * @returns {{ code: string, type: WgslType }}
+ * @returns {{ code: string, type: WgslType, value?: number }}
  */
 function compileExpression(c, node) {
+    const e = compileExpressionAbstract(c, node);
+    return isAbstractFloat(e.type) ? convertAbstract(e, 'f32') : e;
+}
+
+/**
+ * Compile an expression node (any wrapper or specific kind), leaving an abstract result abstract.
+ *
+ * `value` is carried alongside a literal, negated or parenthesised, so that converting it to f32
+ * is a constant fold rather than a call at run time.
+ *
+ * @param {Compiler} c
+ * @param {Node} node
+ * @returns {{ code: string, type: WgslType, value?: number }}
+ */
+function compileExpressionAbstract(c, node) {
     const u = unwrapExpression(node);
 
     // Binary operations
@@ -1809,15 +1861,15 @@ function compileExpression(c, node) {
 
     // Primary subkinds
     if (u.type === 'literal' || u.type === 'int_literal' || u.type === 'float_literal' || u.type === 'bool_literal') {
-        // A bare integer literal is WGSL's AbstractInt. Its concrete type comes from context this
-        // compiler does not track, so the type stays unknown, but operator selection still needs
-        // to know it is an integer: see integerDivision().
-        return { code: literalToJs(u.text), type: literalType(u.text), abstract_int: isBareIntLiteral(u.text) };
+        // A bare literal is abstract (see literalType), and takes its concrete type where it is used
+        const code = literalToJs(u.text);
+        const type = literalType(u.text);
+        return isAbstractType(type) ? { code, type, value: Number(code) } : { code, type };
     }
     if (u.type === 'paren_expression') {
         const inner = namedChildByType(u, 'expression');
-        const e = compileExpression(c, inner);
-        return { code: `(${e.code})`, type: e.type, abstract_int: e.abstract_int };
+        const e = compileExpressionAbstract(c, inner);
+        return { code: `(${e.code})`, type: e.type, value: e.value };
     }
     if (u.type === 'call_expression') {
         const call_phrase = namedChildByType(u, 'call_phrase');
@@ -1849,7 +1901,7 @@ function compileExpression(c, node) {
 function compileBinaryExpression(c, node, op_type, mode) {
     const named = node.namedChildren;
     if (named.length === 1) {
-        return compileExpression(c, named[0]);
+        return compileExpressionAbstract(c, named[0]);
     }
 
     let op_text = null;
@@ -1870,8 +1922,10 @@ function compileBinaryExpression(c, node, op_type, mode) {
         return { code: node.text, type: T_UNKNOWN };
     }
 
-    const left = compileExpression(c, left_node);
-    const right = compileExpression(c, right_node);
+    const [left, right] = concretizeOperands(
+        compileExpressionAbstract(c, left_node),
+        compileExpressionAbstract(c, right_node)
+    );
 
     if (mode === 'arith') {
         const helper = ARITHMETIC_OP[op_text];
@@ -1908,11 +1962,10 @@ function compileBinaryExpression(c, node, op_type, mode) {
             if (f32_scalar !== undefined && isScalarF32(arith_type)) {
                 actual_helper = f32_scalar;
             }
+            // arithmetic on two AbstractInts is still AbstractInt, so `(1 + 2) / 2` is 1
             return {
                 code: `wgsl.${actual_helper}(${left.code}, ${right.code})`,
                 type: arith_type,
-                // arithmetic on two AbstractInts is still AbstractInt: `(1 + 2) / 2` is 1
-                abstract_int: left.abstract_int === true && right.abstract_int === true,
             };
         }
     }
@@ -2009,20 +2062,114 @@ function isIntegerType(t) {
     return false;
 }
 
-/** True if the literal's spelling is a bare integer: no suffix, no fraction, no exponent. */
-function isBareIntLiteral(text) {
-    return /^(0[xX][0-9a-fA-F]+|[0-9]+)$/.test(text.trim());
+/** True if `t` is AbstractInt or AbstractFloat. */
+function isAbstractType(t) {
+    return t !== undefined && t !== null && t.kind === 'abstract';
+}
+
+/** True if `t` is AbstractInt. */
+function isAbstractInt(t) {
+    return isAbstractType(t) && t.name === 'AbstractInt';
+}
+
+/** True if `t` is AbstractFloat. */
+function isAbstractFloat(t) {
+    return isAbstractType(t) && t.name === 'AbstractFloat';
+}
+
+/**
+ * The type a declaration that must be concrete gives an initializer of type `t`: an AbstractInt
+ * becomes i32 and an AbstractFloat f32, as in `let x = 3;` and `let h = 0.5;`. Anything else is
+ * already concrete, or untracked.
+ * @param {WgslType} t
+ * @returns {WgslType}
+ */
+function concreteDefault(t) {
+    if (isAbstractInt(t)) return T_I32;
+    if (isAbstractFloat(t)) return T_F32;
+    return t;
+}
+
+/**
+ * An abstract operand converted to the concrete scalar type it meets.
+ *
+ * Only a float target changes the value. An AbstractFloat becoming an f32 is rounded to one, and
+ * so is an AbstractInt too large for f32 to hold. A literal is folded at compile time, and is
+ * emitted unchanged when f32 holds it already, which most literals in a shader are. An integer
+ * target changes nothing: an AbstractInt has the same value in i32 and u32, and an AbstractFloat
+ * cannot meet an integer in a valid shader.
+ *
+ * @param {{ code: string, type: WgslType, value?: number }} operand
+ * @param {string} scalar the concrete scalar type's name
+ * @returns {{ code: string, type: WgslType }}
+ */
+function convertAbstract(operand, scalar) {
+    if (!isAbstractType(operand.type)) return operand;
+
+    const type = scalar === 'f32' ? T_F32 : { kind: 'primitive', name: scalar };
+
+    if (scalar === 'f32') {
+        if (operand.value === undefined) {
+            return { code: `wgsl.f32(${operand.code})`, type };
+        }
+        const rounded = Math.fround(operand.value);
+        return { code: rounded === operand.value ? operand.code : `(${rounded})`, type };
+    }
+
+    if (scalar === 'f16') {
+        return { code: `wgsl.f16(${operand.code})`, type };
+    }
+
+    return { code: operand.code, type };
+}
+
+/**
+ * The operands of a binary operator, with an abstract one converted to the type of the concrete one
+ * it meets, as WGSL converts it. `A * x` with an f32 x is an f32 product and `A * x` with a u32 x a
+ * wrapping u32 one, for the same `const A = 2;`. Two abstract operands stay abstract: that is a
+ * const-expression, evaluated at full precision.
+ *
+ * An operand the compiler cannot type is almost always a concrete value it cannot see into, like a
+ * struct field or an array element. An AbstractFloat meeting one is taken to meet an f32, since a
+ * float literal can only meet a float and f16 is rare — the type a float literal used to have here
+ * regardless. An AbstractInt stays as it is, since the other side could be any of four types.
+ *
+ * @param {{ code: string, type: WgslType, value?: number }} left
+ * @param {{ code: string, type: WgslType, value?: number }} right
+ * @returns {Array<{ code: string, type: WgslType, value?: number }>}
+ */
+function concretizeOperands(left, right) {
+    const left_abstract = isAbstractType(left.type);
+    const right_abstract = isAbstractType(right.type);
+
+    if (left_abstract === right_abstract) {
+        return [left, right];
+    }
+
+    const abstract = left_abstract ? left : right;
+    const other = left_abstract ? right.type : left.type;
+
+    let scalar = null;
+    if (other && other.kind === 'primitive') scalar = other.name;
+    else if (other && (other.kind === 'vec' || other.kind === 'mat')) scalar = other.scalar;
+    else if (isAbstractFloat(abstract.type)) scalar = 'f32';
+
+    if (scalar === null || scalar === 'bool') {
+        return [left, right];
+    }
+
+    const converted = convertAbstract(abstract, scalar);
+    return left_abstract ? [converted, right] : [left, converted];
 }
 
 /**
  * True if WGSL evaluates `left / right` as integer division. Both operands integer-typed is the
- * obvious case. A bare literal is AbstractInt and takes the other operand's integer type, so
- * `index / 3` truncates; two bare literals divide as AbstractInt, so `7 / 2` is 3 even where an
- * f32 is expected. Negation and parentheses keep a literal abstract: `index / -3` and
- * `index / (3)` truncate too.
+ * obvious case, and an abstract operand has already taken the other's type by now, so `index / 3`
+ * truncates. Two AbstractInts divide as AbstractInt, so `7 / 2` is 3 even where an f32 is expected.
+ * Negation and parentheses keep a literal abstract: `index / -3` and `index / (3)` truncate too.
  */
 function integerDivision(left, right) {
-    const isInt = (operand) => isIntegerType(operand.type) || operand.abstract_int === true;
+    const isInt = (operand) => isIntegerType(operand.type) || isAbstractInt(operand.type);
     return isInt(left) && isInt(right);
 }
 
@@ -2087,7 +2234,8 @@ function scalarTypeSuffix(scalar) {
 
 /**
  * Pick the result type of an element-wise arithmetic op given the operand types.
- * Vec wins over scalar; matrix wins over scalar; scalar otherwise.
+ * Vec wins over scalar; matrix wins over scalar; a concrete scalar over an abstract one. Two
+ * abstract operands give an abstract result, a float one if either is.
  */
 function inferBinaryArithResultType(a, b) {
     const isVec = (t) => t && t.kind === 'vec';
@@ -2100,6 +2248,9 @@ function inferBinaryArithResultType(a, b) {
     if (isMat(b)) return b;
     if (isPrim(a)) return a;
     if (isPrim(b)) return b;
+    if (isAbstractType(a) && isAbstractType(b)) {
+        return isAbstractFloat(a) || isAbstractFloat(b) ? T_ABSTRACT_FLOAT : T_ABSTRACT_INT;
+    }
     return T_UNKNOWN;
 }
 
@@ -2119,6 +2270,8 @@ function inferRawBinaryResultType(a, b, op) {
     if (b && b.kind === 'primitive') return b;
     if (a && a.kind === 'vec') return a;
     if (b && b.kind === 'vec') return b;
+    // `1 << 8` is still an AbstractInt
+    if (isAbstractInt(a) && isAbstractInt(b)) return T_ABSTRACT_INT;
     return T_UNKNOWN;
 }
 
@@ -2127,12 +2280,18 @@ function compileUnaryExpression(c, node) {
     const prefix = textBeforeFirstNamed(node).trim();
 
     if (prefix === '') {
-        return compileExpression(c, named[0]);
+        return compileExpressionAbstract(c, named[0]);
     }
 
-    const operand = compileExpression(c, named[named.length - 1]);
+    const operand = compileExpressionAbstract(c, named[named.length - 1]);
 
-    if (prefix === '-') return { code: `wgsl.neg(${operand.code})`, type: operand.type, abstract_int: operand.abstract_int };
+    if (prefix === '-') {
+        return {
+            code: `wgsl.neg(${operand.code})`,
+            type: operand.type,
+            value: operand.value === undefined ? undefined : -operand.value,
+        };
+    }
     if (prefix === '!') return { code: `!${operand.code}`, type: operand.type };
     if (prefix === '~') return { code: `wgsl.not(${operand.code})`, type: operand.type };
     if (prefix === '*' || prefix === '&') return { code: operand.code, type: operand.type };
@@ -2144,8 +2303,8 @@ function compileSingularExpression(c, node) {
     const named = node.namedChildren;
     const head = named[0];
     const head_result = head.type === 'primary_expression'
-        ? compileExpression(c, unwrapExpression(head))
-        : compileExpression(c, head);
+        ? compileExpressionAbstract(c, unwrapExpression(head))
+        : compileExpressionAbstract(c, head);
 
     let code = head_result.code;
     let type = head_result.type;
@@ -2158,9 +2317,9 @@ function compileSingularExpression(c, node) {
         type = next.type;
     }
 
-    // a bare literal reached through a singular expression is still AbstractInt; anything with an
-    // access specifier on it is not a literal
-    return { code, type, abstract_int: named.length === 1 ? head_result.abstract_int : undefined };
+    // a bare literal reached through a singular expression is still a literal; anything with an
+    // access specifier on it is not one
+    return { code, type, value: named.length === 1 ? head_result.value : undefined };
 }
 
 /**
@@ -2266,15 +2425,27 @@ function compileCallPhraseTyped(c, node) {
     const template = namedChildByType(ident_node, 'template_list');
 
     /** @type {Array<{ code: string, type: WgslType }>} */
-    const arg_results = [];
+    let arg_results = [];
     if (arg_list !== null) {
         const comma_list = namedChildByType(arg_list, 'expression_comma_list');
         if (comma_list !== null) {
             for (const arg of namedChildrenByType(comma_list, 'expression')) {
-                arg_results.push(compileExpression(c, arg));
+                arg_results.push(compileExpressionAbstract(c, arg));
             }
         }
     }
+
+    // A call's arguments are concrete values, as in compileExpression, except for an element-wise
+    // builtin all of whose arguments are abstract: WGSL evaluates `sqrt(2.0)` as an AbstractFloat.
+    const const_evaluated = ELEMENT_WISE_BUILTIN.has(callee_name)
+        && !c.functions.has(callee_name)
+        && arg_results.length > 0
+        && arg_results.every(a => isAbstractType(a.type));
+
+    if (!const_evaluated) {
+        arg_results = arg_results.map(a => isAbstractFloat(a.type) ? convertAbstract(a, 'f32') : a);
+    }
+
     const args = arg_results.map(a => a.code);
 
     // 1) Vec aliases: vec3f, vec4i, vec2h, etc.
@@ -2430,26 +2601,47 @@ function compileCallPhraseTyped(c, node) {
  */
 const INTEGER_TYPED_BUILTIN = new Set(['abs', 'firstLeadingBit', 'firstTrailingBit']);
 
+/**
+ * Element-wise builtins: the result has the type of the arguments. Most are vector-preserving
+ * math; subgroupBroadcastFirst is not math, but has the same shape.
+ * @type {Set<string>}
+ */
+const ELEMENT_WISE_BUILTIN = new Set([
+    'abs', 'sign', 'floor', 'ceil', 'round', 'trunc', 'fract', 'sqrt', 'inverseSqrt',
+    'exp', 'exp2', 'log', 'log2', 'sin', 'cos', 'tan', 'asin', 'acos', 'atan',
+    'pow', 'atan2', 'min', 'max', 'step', 'clamp', 'saturate', 'mix', 'smoothstep',
+    'normalize', 'reflect', 'fma',
+    'subgroupBroadcastFirst',
+]);
+
+/**
+ * The common type of a builtin's same-typed arguments: the first vector or matrix, else the first
+ * concrete scalar — `max(0, x)` with a u32 x is a u32 — else their abstract type if all of them
+ * are abstract, else the first argument's.
+ * @param {Array<{ type: WgslType }>} args
+ * @returns {WgslType}
+ */
+function commonArgumentType(args) {
+    for (const a of args) {
+        if (a.type && (a.type.kind === 'vec' || a.type.kind === 'mat')) return a.type;
+    }
+    for (const a of args) {
+        if (a.type && a.type.kind === 'primitive') return a.type;
+    }
+    if (args.length > 0 && args.every(a => isAbstractType(a.type))) {
+        return args.some(a => isAbstractFloat(a.type)) ? T_ABSTRACT_FLOAT : T_ABSTRACT_INT;
+    }
+    if (args[0] && args[0].type) return args[0].type;
+    return T_UNKNOWN;
+}
+
 /**
  * Best-effort return type for WGSL builtins commonly used in compute shaders.
  * Handles the cases needed for accurate swizzle / member-access codegen downstream.
  */
 function inferBuiltinReturnType(name, args) {
-    // Element-wise vector-preserving math: result follows the first vector argument.
-    const ELEMENT_WISE = new Set([
-        'abs', 'sign', 'floor', 'ceil', 'round', 'trunc', 'fract', 'sqrt', 'inverseSqrt',
-        'exp', 'exp2', 'log', 'log2', 'sin', 'cos', 'tan', 'asin', 'acos', 'atan',
-        'pow', 'atan2', 'min', 'max', 'step', 'clamp', 'saturate', 'mix', 'smoothstep',
-        'normalize', 'reflect', 'fma',
-        // not math, but the same shape: the result has the type of its one argument
-        'subgroupBroadcastFirst',
-    ]);
-    if (ELEMENT_WISE.has(name)) {
-        for (const a of args) {
-            if (a.type && (a.type.kind === 'vec' || a.type.kind === 'mat')) return a.type;
-        }
-        if (args[0] && args[0].type) return args[0].type;
-        return T_UNKNOWN;
+    if (ELEMENT_WISE_BUILTIN.has(name)) {
+        return commonArgumentType(args);
     }
     if (name === 'length' || name === 'distance' || name === 'dot' || name === 'determinant') {
         return { kind: 'primitive', name: 'f32' };
@@ -2461,9 +2653,8 @@ function inferBuiltinReturnType(name, args) {
     }
     if (name === 'all' || name === 'any') return { kind: 'primitive', name: 'bool' };
     if (name === 'select') {
-        if (args[0] && args[0].type) return args[0].type;
-        if (args[1] && args[1].type) return args[1].type;
-        return T_UNKNOWN;
+        // the two values; the third argument is the condition
+        return commonArgumentType(args.slice(0, 2));
     }
     if (name === 'textureLoad' || name === 'textureSampleLevel' || name === 'textureSample' || name === 'textureSampleGrad') {
         return { kind: 'vec', size: 4, scalar: 'f32' };