typegpu/data/generalizeFn.js

A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.

import { vec2b, vec3b, vec4b, vecTypeToConstructor } from "./vector.js";
import { mat2x2f, mat3x3f, mat4x4f } from "./matrix.js";
import { isVecInstance, } from "./wgslTypes.js";
import { invariant } from "../errors.js";
const booleanFor = {
    vec2f: vec2b,
    vec2h: vec2b,
    vec2i: vec2b,
    vec2u: vec2b,
    'vec2<bool>': vec2b,
    vec3f: vec3b,
    vec3h: vec3b,
    vec3i: vec3b,
    vec3u: vec3b,
    'vec3<bool>': vec3b,
    vec4f: vec4b,
    vec4h: vec4b,
    vec4i: vec4b,
    vec4u: vec4b,
    'vec4<bool>': vec4b,
};
const constructorFor = {
    ...vecTypeToConstructor,
    mat2x2f,
    mat3x3f,
    mat4x4f,
};
function getConstructorFor(mode, kind) {
    const map = mode === 'boolean' ? booleanFor : constructorFor;
    if (kind in map) {
        return map[kind];
    }
    throw new Error(`No corresponding vector/matrix type for '${kind}' kind in '${mode}' mode.`);
}
function makeIterable(item) {
    if (item.kind.startsWith('vec')) {
        return item;
    }
    return item.columns.flat();
}
function applyArgs(fn, args, mode) {
    // I'm sorry, TypeScript, I swear I won't lie to you no more ;-;
    const kinds = args.map(kindOf);
    if (kinds.every((type) => type === 'boolean' || type === 'number')) {
        return fn(...args);
    }
    const kind = kinds[0];
    invariant(kind, `Expected kind of the first argument to be present.`);
    const constructor = getConstructorFor(mode, kind);
    const iterableArgs = args.map(makeIterable);
    const length = iterableArgs[0]?.length;
    invariant(length !== undefined, `Expected constructor to have at least one argument.`);
    const constructorArgs = Array.from({ length }, (_, i) => {
        const args = iterableArgs.map((arg) => arg[i]);
        return fn(...args);
    });
    return constructor(...constructorArgs);
}
export function generalizeFn(fn, args) {
    return applyArgs(fn, args, 'first');
}
export function generalizeBoolFn(fn, args) {
    return applyArgs(fn, args, 'boolean');
}
function kindOf(v) {
    if (typeof v === 'number') {
        return 'number';
    }
    if (typeof v === 'boolean') {
        return 'boolean';
    }
    return v.kind;
}
/**
 * If one of the arguments is a vector and other is a number,
 * the number is up-cased to a vector.
 */
export function upCast(args) {
    const [lhs, rhs] = args;
    if (typeof lhs === 'number' && isVecInstance(rhs)) {
        const schema = constructorFor[rhs.kind];
        return [schema(lhs), rhs];
    }
    else if (isVecInstance(lhs) && typeof rhs === 'number') {
        const schema = constructorFor[lhs.kind];
        return [lhs, schema(rhs)];
    }
    return [lhs, rhs];
}