typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
542 lines (541 loc) • 15.8 kB
JavaScript
import { comptime } from "../core/function/comptime.js";
import { callableSchema } from "../core/function/createCallableSchema.js";
import { dualImpl } from "../core/function/dualImpl.js";
import { stitch } from "../core/resolve/stitch.js";
import { $repr } from "../shared/symbols.js";
import { $internal, $resolve } from "../shared/symbols.js";
import { numericLiteralToSnippet } from "../tgsl/generationHelpers.js";
import { f32 } from "./numeric.js";
import {} from "./snippet.js";
import { vec2f, vec3f, vec4f } from "./vector.js";
import { isVec, } from "./wgslTypes.js";
export class MatBase {
}
function createMatSchema(options) {
const construct = callableSchema({
name: options.type,
schema: () => schema,
normalImpl: (...args) => {
const elements = [];
for (const arg of args) {
if (typeof arg === 'number') {
elements.push(arg);
}
else {
for (let i = 0; i < arg.length; ++i) {
elements.push(arg[i]);
}
}
}
if (elements.length !== 0 && elements.length !== options.columns * options.rows) {
throw new Error(`'${options.type}' constructor called with invalid number of arguments.`);
}
for (let i = elements.length; i < options.columns * options.rows; ++i) {
elements.push(0);
}
return new options.MatImpl(...elements);
},
argTypes: (...args) => args.map((arg) => (isVec(arg) ? arg : f32)),
codegenImpl: (ctx, args) => ctx.gen.typeInstantiation(schema, args),
});
const schema = Object.assign(construct, {
[$internal]: {},
type: options.type,
primitive: f32,
identity: identityFunctions[options.columns],
translation: options.columns === 4 ? translation4 : undefined,
scaling: options.columns === 4 ? scaling4 : undefined,
rotationX: options.columns === 4 ? rotationX4 : undefined,
rotationY: options.columns === 4 ? rotationY4 : undefined,
rotationZ: options.columns === 4 ? rotationZ4 : undefined,
});
// TODO: Remove workaround
// it's a workaround for circular dependencies caused by us using schemas in the shader generator
options.MatImpl.prototype.schema = schema;
return schema;
}
const VALID_MAT2x2_ELEMENTS = [0, 1, 2, 3];
class mat2x2Impl extends MatBase {
[$internal] = true;
columns;
length = 4;
constructor(...elements) {
super();
this.columns = [
this.makeColumn(elements[0], elements[1]),
this.makeColumn(elements[2], elements[3]),
];
}
get 0() {
return this.columns[0].x;
}
get 1() {
return this.columns[0].y;
}
get 2() {
return this.columns[1].x;
}
get 3() {
return this.columns[1].y;
}
set 0(value) {
this.columns[0].x = value;
}
set 1(value) {
this.columns[0].y = value;
}
set 2(value) {
this.columns[1].x = value;
}
set 3(value) {
this.columns[1].y = value;
}
*[Symbol.iterator]() {
yield this[0];
yield this[1];
yield this[2];
yield this[3];
}
[$resolve](ctx) {
return ctx.gen.typeInstantiation(mat2x2f,
// oxlint-disable-next-line typescript-eslint(no-non-null-assertion)
VALID_MAT2x2_ELEMENTS.map((i) => numericLiteralToSnippet(this[i])));
}
toString() {
return `${this.kind}(${VALID_MAT2x2_ELEMENTS.map((i) => this[i]).join(', ')})`;
}
}
class mat2x2fImpl extends mat2x2Impl {
kind = 'mat2x2f';
makeColumn(e0, e1) {
return vec2f(e0, e1);
}
}
const VALID_MAT3x3_ELEMENTS = [0, 1, 2, 4, 5, 6, 8, 9, 10];
class mat3x3Impl extends MatBase {
[$internal] = true;
columns;
length = 12;
constructor(...elements) {
super();
this.columns = [
this.makeColumn(elements[0], elements[1], elements[2]),
this.makeColumn(elements[3], elements[4], elements[5]),
this.makeColumn(elements[6], elements[7], elements[8]),
];
}
get 0() {
return this.columns[0].x;
}
get 1() {
return this.columns[0].y;
}
get 2() {
return this.columns[0].z;
}
get 3() {
return 0;
}
get 4() {
return this.columns[1].x;
}
get 5() {
return this.columns[1].y;
}
get 6() {
return this.columns[1].z;
}
get 7() {
return 0;
}
get 8() {
return this.columns[2].x;
}
get 9() {
return this.columns[2].y;
}
get 10() {
return this.columns[2].z;
}
get 11() {
return 0;
}
set 0(value) {
this.columns[0].x = value;
}
set 1(value) {
this.columns[0].y = value;
}
set 2(value) {
this.columns[0].z = value;
}
set 3(_) { }
set 4(value) {
this.columns[1].x = value;
}
set 5(value) {
this.columns[1].y = value;
}
set 6(value) {
this.columns[1].z = value;
}
set 7(_) { }
set 8(value) {
this.columns[2].x = value;
}
set 9(value) {
this.columns[2].y = value;
}
set 10(value) {
this.columns[2].z = value;
}
set 11(_) { }
*[Symbol.iterator]() {
for (let i = 0; i < 12; i++) {
yield this[i];
}
}
[$resolve](ctx) {
return ctx.gen.typeInstantiation(mat3x3f,
// oxlint-disable-next-line typescript-eslint(no-non-null-assertion)
VALID_MAT3x3_ELEMENTS.map((i) => numericLiteralToSnippet(this[i])));
}
toString() {
return `${this.kind}(${VALID_MAT3x3_ELEMENTS.map((i) => this[i]).join(', ')})`;
}
}
class mat3x3fImpl extends mat3x3Impl {
kind = 'mat3x3f';
makeColumn(x, y, z) {
return vec3f(x, y, z);
}
}
const VALID_MAT4x4_ELEMENTS = Array.from({ length: 16 }, (_, i) => i);
class mat4x4Impl extends MatBase {
[$internal] = true;
columns;
constructor(...elements) {
super();
this.columns = [
this.makeColumn(elements[0], elements[1], elements[2], elements[3]),
this.makeColumn(elements[4], elements[5], elements[6], elements[7]),
this.makeColumn(elements[8], elements[9], elements[10], elements[11]),
this.makeColumn(elements[12], elements[13], elements[14], elements[15]),
];
}
length = 16;
get 0() {
return this.columns[0].x;
}
get 1() {
return this.columns[0].y;
}
get 2() {
return this.columns[0].z;
}
get 3() {
return this.columns[0].w;
}
get 4() {
return this.columns[1].x;
}
get 5() {
return this.columns[1].y;
}
get 6() {
return this.columns[1].z;
}
get 7() {
return this.columns[1].w;
}
get 8() {
return this.columns[2].x;
}
get 9() {
return this.columns[2].y;
}
get 10() {
return this.columns[2].z;
}
get 11() {
return this.columns[2].w;
}
get 12() {
return this.columns[3].x;
}
get 13() {
return this.columns[3].y;
}
get 14() {
return this.columns[3].z;
}
get 15() {
return this.columns[3].w;
}
set 0(value) {
this.columns[0].x = value;
}
set 1(value) {
this.columns[0].y = value;
}
set 2(value) {
this.columns[0].z = value;
}
set 3(value) {
this.columns[0].w = value;
}
set 4(value) {
this.columns[1].x = value;
}
set 5(value) {
this.columns[1].y = value;
}
set 6(value) {
this.columns[1].z = value;
}
set 7(value) {
this.columns[1].w = value;
}
set 8(value) {
this.columns[2].x = value;
}
set 9(value) {
this.columns[2].y = value;
}
set 10(value) {
this.columns[2].z = value;
}
set 11(value) {
this.columns[2].w = value;
}
set 12(value) {
this.columns[3].x = value;
}
set 13(value) {
this.columns[3].y = value;
}
set 14(value) {
this.columns[3].z = value;
}
set 15(value) {
this.columns[3].w = value;
}
*[Symbol.iterator]() {
for (let i = 0; i < 16; i++) {
yield this[i];
}
}
[$resolve](ctx) {
return ctx.gen.typeInstantiation(mat4x4f,
// oxlint-disable-next-line typescript-eslint(no-non-null-assertion)
VALID_MAT4x4_ELEMENTS.map((i) => numericLiteralToSnippet(this[i])));
}
toString() {
return `${this.kind}(${VALID_MAT4x4_ELEMENTS.map((i) => this[i]).join(', ')})`;
}
}
class mat4x4fImpl extends mat4x4Impl {
kind = 'mat4x4f';
makeColumn(x, y, z, w) {
return vec4f(x, y, z, w);
}
}
// ----------
// Matrix ops
// ----------
/**
* Returns a 2-by-2 identity matrix.
* @returns {m2x2f} The result matrix.
*/
export const identity2 = comptime(() => mat2x2f(1, 0, 0, 1)).$name('identity2');
/**
* Returns a 3-by-3 identity matrix.
* @returns {m3x3f} The result matrix.
*/
export const identity3 = comptime(() => mat3x3f(1, 0, 0, 0, 1, 0, 0, 0, 1)).$name('identity3');
/**
* Returns a 4-by-4 identity matrix.
* @returns {m4x4f} The result matrix.
*/
export const identity4 = comptime(() => mat4x4f(1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1)).$name('identity4');
const identityFunctions = {
2: identity2,
3: identity3,
4: identity4,
};
/**
* Creates a 4-by-4 matrix which translates by the given vector v.
* @param {v3f} vector - The vector by which to translate.
* @returns {m4x4f} The translation matrix.
*/
export const translation4 = dualImpl({
name: 'translation4',
normalImpl: (vector) =>
// oxfmt-ignore
mat4x4f(1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, vector.x, vector.y, vector.z, 1),
get signature() {
return { argTypes: [vec3f], returnType: mat4x4f };
},
codegenImpl: (_ctx, [v]) => stitch `mat4x4f(1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, ${v}.x, ${v}.y, ${v}.z, 1)`,
sideEffects: false,
});
/**
* Creates a 4-by-4 matrix which scales in each dimension by an amount given by the corresponding entry in the given vector.
* @param {v3f} vector - A vector of three entries specifying the factor by which to scale in each dimension.
* @returns {m4x4f} The scaling matrix.
*/
export const scaling4 = dualImpl({
name: 'scaling4',
normalImpl: (vector) =>
// oxfmt-ignore
mat4x4f(vector.x, 0, 0, 0, 0, vector.y, 0, 0, 0, 0, vector.z, 0, 0, 0, 0, 1),
get signature() {
return { argTypes: [vec3f], returnType: mat4x4f };
},
codegenImpl: (_ctx, [v]) => stitch `mat4x4f(${v}.x, 0, 0, 0, 0, ${v}.y, 0, 0, 0, 0, ${v}.z, 0, 0, 0, 0, 1)`,
sideEffects: false,
});
/**
* Creates a 4-by-4 matrix which rotates around the x-axis by the given angle.
* @param {number} angle - The angle by which to rotate (in radians).
* @returns {m4x4f} The rotation matrix.
*/
export const rotationX4 = dualImpl({
name: 'rotationX4',
normalImpl: (a) =>
// oxfmt-ignore
mat4x4f(1, 0, 0, 0, 0, Math.cos(a), Math.sin(a), 0, 0, -Math.sin(a), Math.cos(a), 0, 0, 0, 0, 1),
get signature() {
return { argTypes: [f32], returnType: mat4x4f };
},
codegenImpl: (_ctx, [a]) => stitch `mat4x4f(1, 0, 0, 0, 0, cos(${a}), sin(${a}), 0, 0, -sin(${a}), cos(${a}), 0, 0, 0, 0, 1)`,
sideEffects: false,
});
/**
* Creates a 4-by-4 matrix which rotates around the y-axis by the given angle.
* @param {number} angle - The angle by which to rotate (in radians).
* @returns {m4x4f} The rotation matrix.
*/
export const rotationY4 = dualImpl({
name: 'rotationY4',
normalImpl: (a) =>
// oxfmt-ignore
mat4x4f(Math.cos(a), 0, -Math.sin(a), 0, 0, 1, 0, 0, Math.sin(a), 0, Math.cos(a), 0, 0, 0, 0, 1),
get signature() {
return { argTypes: [f32], returnType: mat4x4f };
},
codegenImpl: (_ctx, [a]) => stitch `mat4x4f(cos(${a}), 0, -sin(${a}), 0, 0, 1, 0, 0, sin(${a}), 0, cos(${a}), 0, 0, 0, 0, 1)`,
sideEffects: false,
});
/**
* Creates a 4-by-4 matrix which rotates around the z-axis by the given angle.
* @param {number} angle - The angle by which to rotate (in radians).
* @returns {m4x4f} The rotation matrix.
*/
export const rotationZ4 = dualImpl({
name: 'rotationZ4',
normalImpl: (a) =>
// oxfmt-ignore
mat4x4f(Math.cos(a), Math.sin(a), 0, 0, -Math.sin(a), Math.cos(a), 0, 0, 0, 0, 1, 0, 0, 0, 0, 1),
get signature() {
return { argTypes: [f32], returnType: mat4x4f };
},
codegenImpl: (_ctx, [a]) => stitch `mat4x4f(cos(${a}), sin(${a}), 0, 0, -sin(${a}), cos(${a}), 0, 0, 0, 0, 1, 0, 0, 0, 0, 1)`,
sideEffects: false,
});
// ----------
// Public API
// ----------
/**
* Schema representing mat2x2f - a matrix with 2 rows and 2 columns, with elements of type f32.
* Also a constructor function for this matrix type.
*
* @example
* const zero2x2 = mat2x2f(); // filled with zeros
*
* @example
* const mat = mat2x2f(0, 1, 2, 3);
* mat.columns[0] // vec2f(0, 1)
* mat.columns[1] // vec2f(2, 3)
*
* @example
* const mat = mat2x2f(
* vec2f(0, 1), // column 0
* vec2f(1, 2), // column 1
* );
*
* @example
* const buffer = root.createBuffer(d.mat2x2f, d.mat2x2f(0, 1, 2, 3)); // buffer holding a d.mat2x2f value, with an initial value of ((0, 1), (2, 3))
*/
export const mat2x2f = createMatSchema({
type: 'mat2x2f',
rows: 2,
columns: 2,
MatImpl: mat2x2fImpl,
});
/**
* Schema representing mat3x3f - a matrix with 3 rows and 3 columns, with elements of type f32.
* Also a constructor function for this matrix type.
*
* @example
* const zero3x3 = mat3x3f(); // filled with zeros
*
* @example
* const mat = mat3x3f(0, 1, 2, 3, 4, 5, 6, 7, 8);
* mat.columns[0] // vec3f(0, 1, 2)
* mat.columns[1] // vec3f(3, 4, 5)
* mat.columns[2] // vec3f(6, 7, 8)
*
* @example
* const mat = mat3x3f(
* vec3f(0, 1, 2), // column 0
* vec3f(2, 3, 4), // column 1
* vec3f(5, 6, 7), // column 2
* );
*
* @example
* const buffer = root.createBuffer(d.mat3x3f, d.mat3x3f()); // buffer holding a d.mat3x3f value, with an initial value of mat3x3f filled with zeros
*/
export const mat3x3f = createMatSchema({
type: 'mat3x3f',
rows: 3,
columns: 3,
MatImpl: mat3x3fImpl,
});
/**
* Schema representing mat4x4f - a matrix with 4 rows and 4 columns, with elements of type f32.
* Also a constructor function for this matrix type.
*
* @example
* const zero4x4 = mat4x4f(); // filled with zeros
*
* @example
* const mat = mat4x4f(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
* mat.columns[0] // vec4f(0, 1, 2, 3)
* mat.columns[1] // vec4f(4, 5, 6, 7)
* mat.columns[2] // vec4f(8, 9, 10, 11)
* mat.columns[3] // vec4f(12, 13, 14, 15)
*
* @example
* const mat = mat4x4f(
* vec4f(0, 1, 2, 3), // column 0
* vec4f(4, 5, 6, 7), // column 1
* vec4f(8, 9, 10, 11), // column 2
* vec4f(12, 13, 14, 15), // column 3
* );
*
* @example
* const buffer = root.createBuffer(d.mat4x4f, d.mat4x4f()); // buffer holding a d.mat4x4f value, with an initial value of mat4x4f filled with zeros
*/
export const mat4x4f = createMatSchema({
type: 'mat4x4f',
rows: 4,
columns: 4,
MatImpl: mat4x4fImpl,
});
export function matToArray(mat) {
if (mat.kind === 'mat3x3f') {
return [mat[0], mat[1], mat[2], mat[4], mat[5], mat[6], mat[8], mat[9], mat[10]];
}
return Array.from({ length: mat.length }).map((_, idx) => mat[idx]);
}