typegpu/data/vectorOps.js

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

import { mat2x2f, mat3x3f, mat4x4f } from "./matrix.js";
import { bitcastF32toU32Impl, bitcastU32toF32Impl, bitcastU32toI32Impl } from "./numberOps.js";
import { vec2f, vec2i, vec2u, vec3f, vec3h, vec3i, vec3u, vec4f, vec4i, vec4u } from "./vector.js";
const lengthVec2 = (v) => Math.sqrt(v.x ** 2 + v.y ** 2);
const lengthVec3 = (v) => Math.sqrt(v.x ** 2 + v.y ** 2 + v.z ** 2);
const lengthVec4 = (v) => Math.sqrt(v.x ** 2 + v.y ** 2 + v.z ** 2 + v.w ** 2);
const dotVec2 = (lhs, rhs) => lhs.x * rhs.x + lhs.y * rhs.y;
const dotVec3 = (lhs, rhs) => lhs.x * rhs.x + lhs.y * rhs.y + lhs.z * rhs.z;
const dotVec4 = (lhs, rhs) => lhs.x * rhs.x + lhs.y * rhs.y + lhs.z * rhs.z + lhs.w * rhs.w;
const binaryComponentWise2u = (op) => (a, b) => vec2u(op(a.x, b.x), op(a.y, b.y));
const binaryComponentWise3u = (op) => (a, b) => vec3u(op(a.x, b.x), op(a.y, b.y), op(a.z, b.z));
const binaryComponentWise4u = (op) => (a, b) => vec4u(op(a.x, b.x), op(a.y, b.y), op(a.z, b.z), op(a.w, b.w));
const binaryComponentWise2i2u = (op) => (a, b) => vec2i(op(a.x, b.x), op(a.y, b.y));
const binaryComponentWise3i3u = (op) => (a, b) => vec3i(op(a.x, b.x), op(a.y, b.y), op(a.z, b.z));
const binaryComponentWise4i4u = (op) => (a, b) => vec4i(op(a.x, b.x), op(a.y, b.y), op(a.z, b.z), op(a.w, b.w));
/**
 * Functions that cannot be simply generalized via `generalizeFn`
 * have their overloads listed explicitly here.
 */
export const VectorOps = {
    all: {
        'vec2<bool>': (e) => e.x && e.y,
        'vec3<bool>': (e) => e.x && e.y && e.z,
        'vec4<bool>': (e) => e.x && e.y && e.z && e.w,
    },
    length: {
        vec2f: lengthVec2,
        vec2h: lengthVec2,
        vec3f: lengthVec3,
        vec3h: lengthVec3,
        vec4f: lengthVec4,
        vec4h: lengthVec4,
    },
    mulMxM: {
        mat2x2f: (a, b) => {
            const a_ = a.columns;
            const b_ = b.columns;
            return mat2x2f(a_[0].x * b_[0].x + a_[1].x * b_[0].y, a_[0].y * b_[0].x + a_[1].y * b_[0].y, a_[0].x * b_[1].x + a_[1].x * b_[1].y, a_[0].y * b_[1].x + a_[1].y * b_[1].y);
        },
        mat3x3f: (a, b) => {
            const a_ = a.columns;
            const b_ = b.columns;
            return mat3x3f(a_[0].x * b_[0].x + a_[1].x * b_[0].y + a_[2].x * b_[0].z, a_[0].y * b_[0].x + a_[1].y * b_[0].y + a_[2].y * b_[0].z, a_[0].z * b_[0].x + a_[1].z * b_[0].y + a_[2].z * b_[0].z, a_[0].x * b_[1].x + a_[1].x * b_[1].y + a_[2].x * b_[1].z, a_[0].y * b_[1].x + a_[1].y * b_[1].y + a_[2].y * b_[1].z, a_[0].z * b_[1].x + a_[1].z * b_[1].y + a_[2].z * b_[1].z, a_[0].x * b_[2].x + a_[1].x * b_[2].y + a_[2].x * b_[2].z, a_[0].y * b_[2].x + a_[1].y * b_[2].y + a_[2].y * b_[2].z, a_[0].z * b_[2].x + a_[1].z * b_[2].y + a_[2].z * b_[2].z);
        },
        mat4x4f: (a, b) => {
            const a_ = a.columns;
            const b_ = b.columns;
            return mat4x4f(a_[0].x * b_[0].x + a_[1].x * b_[0].y + a_[2].x * b_[0].z + a_[3].x * b_[0].w, a_[0].y * b_[0].x + a_[1].y * b_[0].y + a_[2].y * b_[0].z + a_[3].y * b_[0].w, a_[0].z * b_[0].x + a_[1].z * b_[0].y + a_[2].z * b_[0].z + a_[3].z * b_[0].w, a_[0].w * b_[0].x + a_[1].w * b_[0].y + a_[2].w * b_[0].z + a_[3].w * b_[0].w, a_[0].x * b_[1].x + a_[1].x * b_[1].y + a_[2].x * b_[1].z + a_[3].x * b_[1].w, a_[0].y * b_[1].x + a_[1].y * b_[1].y + a_[2].y * b_[1].z + a_[3].y * b_[1].w, a_[0].z * b_[1].x + a_[1].z * b_[1].y + a_[2].z * b_[1].z + a_[3].z * b_[1].w, a_[0].w * b_[1].x + a_[1].w * b_[1].y + a_[2].w * b_[1].z + a_[3].w * b_[1].w, a_[0].x * b_[2].x + a_[1].x * b_[2].y + a_[2].x * b_[2].z + a_[3].x * b_[2].w, a_[0].y * b_[2].x + a_[1].y * b_[2].y + a_[2].y * b_[2].z + a_[3].y * b_[2].w, a_[0].z * b_[2].x + a_[1].z * b_[2].y + a_[2].z * b_[2].z + a_[3].z * b_[2].w, a_[0].w * b_[2].x + a_[1].w * b_[2].y + a_[2].w * b_[2].z + a_[3].w * b_[2].w, a_[0].x * b_[3].x + a_[1].x * b_[3].y + a_[2].x * b_[3].z + a_[3].x * b_[3].w, a_[0].y * b_[3].x + a_[1].y * b_[3].y + a_[2].y * b_[3].z + a_[3].y * b_[3].w, a_[0].z * b_[3].x + a_[1].z * b_[3].y + a_[2].z * b_[3].z + a_[3].z * b_[3].w, a_[0].w * b_[3].x + a_[1].w * b_[3].y + a_[2].w * b_[3].z + a_[3].w * b_[3].w);
        },
    },
    mulMxV: {
        mat2x2f: (m, v) => {
            const m_ = m.columns;
            return vec2f(m_[0].x * v.x + m_[1].x * v.y, m_[0].y * v.x + m_[1].y * v.y);
        },
        mat3x3f: (m, v) => {
            const m_ = m.columns;
            return vec3f(m_[0].x * v.x + m_[1].x * v.y + m_[2].x * v.z, m_[0].y * v.x + m_[1].y * v.y + m_[2].y * v.z, m_[0].z * v.x + m_[1].z * v.y + m_[2].z * v.z);
        },
        mat4x4f: (m, v) => {
            const m_ = m.columns;
            return vec4f(m_[0].x * v.x + m_[1].x * v.y + m_[2].x * v.z + m_[3].x * v.w, m_[0].y * v.x + m_[1].y * v.y + m_[2].y * v.z + m_[3].y * v.w, m_[0].z * v.x + m_[1].z * v.y + m_[2].z * v.z + m_[3].z * v.w, m_[0].w * v.x + m_[1].w * v.y + m_[2].w * v.z + m_[3].w * v.w);
        },
    },
    mulVxM: {
        mat2x2f: (v, m) => {
            const m_ = m.columns;
            return vec2f(v.x * m_[0].x + v.y * m_[0].y, v.x * m_[1].x + v.y * m_[1].y);
        },
        mat3x3f: (v, m) => {
            const m_ = m.columns;
            return vec3f(v.x * m_[0].x + v.y * m_[0].y + v.z * m_[0].z, v.x * m_[1].x + v.y * m_[1].y + v.z * m_[1].z, v.x * m_[2].x + v.y * m_[2].y + v.z * m_[2].z);
        },
        mat4x4f: (v, m) => {
            const m_ = m.columns;
            return vec4f(v.x * m_[0].x + v.y * m_[0].y + v.z * m_[0].z + v.w * m_[0].w, v.x * m_[1].x + v.y * m_[1].y + v.z * m_[1].z + v.w * m_[1].w, v.x * m_[2].x + v.y * m_[2].y + v.z * m_[2].z + v.w * m_[2].w, v.x * m_[3].x + v.y * m_[3].y + v.z * m_[3].z + v.w * m_[3].w);
        },
    },
    dot: {
        vec2f: dotVec2,
        vec2h: dotVec2,
        vec2i: dotVec2,
        vec2u: dotVec2,
        vec3f: dotVec3,
        vec3h: dotVec3,
        vec3i: dotVec3,
        vec3u: dotVec3,
        vec4f: dotVec4,
        vec4h: dotVec4,
        vec4i: dotVec4,
        vec4u: dotVec4,
    },
    cross: {
        vec3f: (a, b) => {
            return vec3f(a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x);
        },
        vec3h: (a, b) => {
            return vec3h(a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x);
        },
    },
    bitShiftLeft: {
        vec2i: binaryComponentWise2i2u((a, b) => a << b),
        vec2u: binaryComponentWise2u((a, b) => a << b),
        vec3i: binaryComponentWise3i3u((a, b) => a << b),
        vec3u: binaryComponentWise3u((a, b) => a << b),
        vec4i: binaryComponentWise4i4u((a, b) => a << b),
        vec4u: binaryComponentWise4u((a, b) => a << b),
    },
    bitShiftRight: {
        vec2i: binaryComponentWise2i2u((a, b) => a >> b),
        vec2u: binaryComponentWise2u((a, b) => a >>> b),
        vec3i: binaryComponentWise3i3u((a, b) => a >> b),
        vec3u: binaryComponentWise3u((a, b) => a >>> b),
        vec4i: binaryComponentWise4i4u((a, b) => a >> b),
        vec4u: binaryComponentWise4u((a, b) => a >>> b),
    },
    bitcastU32toF32: {
        vec2u: (n) => vec2f(bitcastU32toF32Impl(n.x), bitcastU32toF32Impl(n.y)),
        vec3u: (n) => vec3f(bitcastU32toF32Impl(n.x), bitcastU32toF32Impl(n.y), bitcastU32toF32Impl(n.z)),
        vec4u: (n) => vec4f(bitcastU32toF32Impl(n.x), bitcastU32toF32Impl(n.y), bitcastU32toF32Impl(n.z), bitcastU32toF32Impl(n.w)),
    },
    bitcastU32toI32: {
        vec2u: (n) => vec2i(bitcastU32toI32Impl(n.x), bitcastU32toI32Impl(n.y)),
        vec3u: (n) => vec3i(bitcastU32toI32Impl(n.x), bitcastU32toI32Impl(n.y), bitcastU32toI32Impl(n.z)),
        vec4u: (n) => vec4i(bitcastU32toI32Impl(n.x), bitcastU32toI32Impl(n.y), bitcastU32toI32Impl(n.z), bitcastU32toI32Impl(n.w)),
    },
    bitcastF32toU32: {
        vec2f: (n) => vec2u(bitcastF32toU32Impl(n.x), bitcastF32toU32Impl(n.y)),
        vec3f: (n) => vec3u(bitcastF32toU32Impl(n.x), bitcastF32toU32Impl(n.y), bitcastF32toU32Impl(n.z)),
        vec4f: (n) => vec4u(bitcastF32toU32Impl(n.x), bitcastF32toU32Impl(n.y), bitcastF32toU32Impl(n.z), bitcastF32toU32Impl(n.w)),
    },
};