diff --git a/src/strands/ir_builders.js b/src/strands/ir_builders.js index 0c4b139d75..c95c9e07cf 100644 --- a/src/strands/ir_builders.js +++ b/src/strands/ir_builders.js @@ -72,6 +72,56 @@ export function unaryOpNode(strandsContext, nodeOrValue, opCode) { return { id, dimension: node.dimension }; } +function isMatrixArithmeticOp(opCode) { + return ( + opCode === OpCode.Binary.ADD || + opCode === OpCode.Binary.SUBTRACT || + opCode === OpCode.Binary.MULTIPLY || + opCode === OpCode.Binary.DIVIDE + ); +} + +function resolveMatrixBinaryOpType(leftType, rightType, opCode) { + const leftIsMat = leftType.baseType === BaseType.MAT; + const rightIsMat = rightType.baseType === BaseType.MAT; + const matType = leftIsMat ? leftType : rightType; + const leftIsVector = !leftIsMat && leftType.dimension > 1; + const rightIsVector = !rightIsMat && rightType.dimension > 1; + + if (opCode === OpCode.Binary.MULTIPLY) { + if (leftIsMat && rightIsMat) { + if (leftType.dimension !== rightType.dimension) { + FES.userError('type error', + `You can only multiply two matrices of the same size, but got mat${leftType.dimension} * mat${rightType.dimension}.`); + } + return { baseType: BaseType.MAT, dimension: matType.dimension }; + } + if ((leftIsMat && rightIsVector) || (rightIsMat && leftIsVector)) { + const vectorType = leftIsMat ? rightType : leftType; + if (vectorType.dimension !== matType.dimension) { + FES.userError('type error', + `A mat${matType.dimension} can only be multiplied with a vector of length ${matType.dimension}, but got a vector of length ${vectorType.dimension}.`); + } + return { baseType: BaseType.FLOAT, dimension: matType.dimension }; + } + return { baseType: BaseType.MAT, dimension: matType.dimension }; + } + + if (leftIsMat && rightIsMat) { + if (leftType.dimension !== rightType.dimension) { + FES.userError('type error', + `You can only combine two matrices of the same size, but got mat${leftType.dimension} ${OpCodeToSymbol[opCode]} mat${rightType.dimension}.`); + } + return { baseType: BaseType.MAT, dimension: matType.dimension }; + } + if ((leftIsMat && !rightIsVector) || (rightIsMat && !leftIsVector)) { + return { baseType: BaseType.MAT, dimension: matType.dimension }; + } + FES.userError('type error', + `A matrix can't be combined with a vector using '${OpCodeToSymbol[opCode]}'. ` + + `Use matrix multiplication (*) to transform a vector by a matrix.`); +} + export function binaryOpNode(strandsContext, leftStrandsNode, rightArg, opCode) { const { dag, cfg } = strandsContext; // Construct a node for right if its just an array or number etc. @@ -97,6 +147,24 @@ export function binaryOpNode(strandsContext, leftStrandsNode, rightArg, opCode) DAG.propagateTypeToAssignOnUse(dag, rightStrandsNode.id, leftType.baseType, leftType.dimension); rightType = DAG.extractNodeTypeInfo(dag, rightStrandsNode.id); } + + if ( + (leftType.baseType === BaseType.MAT || rightType.baseType === BaseType.MAT) && + isMatrixArithmeticOp(opCode) + ) { + const resultType = resolveMatrixBinaryOpType(leftType, rightType, opCode); + const nodeData = DAG.createNodeData({ + nodeType: NodeType.OPERATION, + opCode, + dependsOn: [finalLeftNodeID, finalRightNodeID], + baseType: resultType.baseType, + dimension: resultType.dimension, + }); + const id = DAG.getOrCreateNode(dag, nodeData); + CFG.recordInBasicBlock(cfg, cfg.currentBlock, id); + return { id, dimension: nodeData.dimension }; + } + const cast = { node: null, toType: leftType }; const bothDeferred = leftType.baseType === rightType.baseType && leftType.baseType === BaseType.DEFER; if (bothDeferred) { @@ -290,6 +358,157 @@ export function constructTypeFromIDs(strandsContext, typeInfo, strandsNodesArray return id; } +export function matrixConstructorNode(strandsContext, dimension, values) { + const { cfg } = strandsContext; + const componentIDs = values.map((value) => { + if (value?.isStrandsNode) { + return value.id; + } + const { id } = scalarLiteralNode( + strandsContext, + { dimension: 1, baseType: BaseType.FLOAT }, + value + ); + return id; + }); + const id = constructTypeFromIDs( + strandsContext, + { baseType: BaseType.MAT, dimension }, + componentIDs + ); + CFG.recordInBasicBlock(cfg, cfg.currentBlock, id); + return { id, dimension }; +} + +export function diagonalMatrixNode(strandsContext, dimension, value) { + const values = []; + for (let col = 0; col < dimension; col++) { + for (let row = 0; row < dimension; row++) { + values.push(row === col ? value : 0); + } + } + return matrixConstructorNode(strandsContext, dimension, values); +} + +export function identityMatrixNode(strandsContext, dimension) { + return diagonalMatrixNode(strandsContext, dimension, 1); +} + +// Convert a matrix into one of a different size, mirroring GLSL's mat3(m4) / +// mat4(m3). Truncation keeps the upper-left submatrix; extension pads the extra +// rows and columns with the identity matrix. This is built from column access + +// swizzles so the GLSL and WGSL backends emit the same code (WGSL has no +// matrix-resizing constructor of its own). +export function matrixResizeNode(strandsContext, srcNode, srcDim, dstDim) { + const { dag, cfg } = strandsContext; + + const sourceColumn = (colIndex) => { + const { id: indexID } = scalarLiteralNode( + strandsContext, + { dimension: 1, baseType: BaseType.INT }, + colIndex + ); + const nodeData = DAG.createNodeData({ + nodeType: NodeType.OPERATION, + opCode: OpCode.Binary.ARRAY_ACCESS, + dependsOn: [srcNode.id, indexID], + dimension: srcDim, + baseType: BaseType.FLOAT, + }); + const id = DAG.getOrCreateNode(dag, nodeData); + CFG.recordInBasicBlock(cfg, cfg.currentBlock, id); + return createStrandsNode(id, srcDim, strandsContext); + }; + + const columns = []; + for (let col = 0; col < dstDim; col++) { + if (col < srcDim && dstDim < srcDim) { + // Truncate: keep the first dstDim components of the source column. + const { id, dimension } = swizzleNode( + strandsContext, sourceColumn(col), 'xyzw'.slice(0, dstDim) + ); + columns.push(createStrandsNode(id, dimension, strandsContext)); + } else if (col < srcDim) { + // Extend: keep the source column, padding the extra rows with zeros. + const pad = Array(dstDim - srcDim).fill(0); + const { id, dimension } = primitiveConstructorNode( + strandsContext, + { baseType: BaseType.FLOAT, dimension: dstDim }, + [sourceColumn(col), ...pad] + ); + columns.push(createStrandsNode(id, dimension, strandsContext)); + } else { + // Identity column for indices beyond the source matrix. + const comps = Array.from({ length: dstDim }, (_, row) => (row === col ? 1 : 0)); + const { id, dimension } = primitiveConstructorNode( + strandsContext, + { baseType: BaseType.FLOAT, dimension: dstDim }, + comps + ); + columns.push(createStrandsNode(id, dimension, strandsContext)); + } + } + + return matrixConstructorNode(strandsContext, dstDim, columns); +} + +export function matrixNode(strandsContext, dimension, args) { + const dag = strandsContext.dag; + const componentCount = dimension * dimension; + const baseTypeOf = (a) => DAG.extractNodeTypeInfo(dag, a.id).baseType; + + // No arguments -> identity. + if (args.length === 0) { + return identityMatrixNode(strandsContext, dimension); + } + + if (args.length === 1) { + const arg = args[0]; + const isNode = !!arg?.isStrandsNode; + + // A matrix of the same size -> copy it. + if (isNode && baseTypeOf(arg) === BaseType.MAT && arg.dimension === dimension) { + return { id: arg.id, dimension }; + } + + // A matrix of a different size -> resize conversion. mat3(someMat4) keeps the + // upper-left 3x3, and mat4(someMat3) extends it with the identity matrix. + // WGSL has no matrix-resizing constructor, so we build the result from + // explicit column swizzles at the IR level, which both backends emit + // identically (e.g. m[0].xyz). + if (isNode && baseTypeOf(arg) === BaseType.MAT) { + return matrixResizeNode(strandsContext, arg, arg.dimension, dimension); + } + + // A single scalar -> diagonal matrix (GLSL's matN(s) idiom, e.g. mat4(1.0)). + const isScalar = typeof arg === 'number' || + (isNode && arg.dimension === 1 && baseTypeOf(arg) !== BaseType.MAT); + if (isScalar) { + return diagonalMatrixNode(strandsContext, dimension, arg); + } + } + + // N column vectors, each of length N. + if ( + args.length === dimension && + args.every((a) => a?.isStrandsNode && a.dimension === dimension && baseTypeOf(a) !== BaseType.MAT) + ) { + return matrixConstructorNode(strandsContext, dimension, args); + } + + // N*N individual components (numbers or scalar nodes), column-major. + if (args.length === componentCount) { + return matrixConstructorNode(strandsContext, dimension, args); + } + + FES.userError( + 'parameter validation', + `A mat${dimension} constructor expects no arguments (identity), a single number ` + + `(diagonal), a single mat${dimension}, ${dimension} column vectors, or ` + + `${componentCount} values — but got ${args.length} argument(s).` + ); +} + export function primitiveConstructorNode(strandsContext, typeInfo, dependsOn) { const cfg = strandsContext.cfg; dependsOn = (Array.isArray(dependsOn) ? dependsOn : [dependsOn]) diff --git a/src/strands/strands_api.js b/src/strands/strands_api.js index e6716e1c9b..4483d4f158 100644 --- a/src/strands/strands_api.js +++ b/src/strands/strands_api.js @@ -837,6 +837,204 @@ const _rgb2hsl = (instance, colorNode) => { }); } + const matrixConstructors = [ + // 2x2 has no affine/transform form, so it only has matrix names. + ['mat2', 2], ['mat2x2', 2], + // 3x3 = 2D affine transform. + ['transform2D', 3], ['mat3', 3], ['mat3x3', 3], + // 4x4 = 3D affine transform. + ['transform3D', 4], ['mat4', 4], ['mat4x4', 4], + ]; + for (const [name, dimension] of matrixConstructors) { + augmentFn(fn, p5, name, function (...args) { + if (!strandsContext.active) { + p5._friendlyError( + `It looks like you've called ${name} outside of a shader's modify() function.` + ); + return; + } + const { id, dimension: dim } = build.matrixNode(strandsContext, dimension, args); + return createStrandsNode(id, dim, strandsContext); + }); + } + + + const isStrandsTransform = (t) => + strandsContext.active && !!t?.isStrandsNode && t.typeInfo().baseType === BaseType.MAT; + + const transformStep = (t, values2D, values3D) => { + if (!t?.isStrandsNode || t.typeInfo().baseType !== BaseType.MAT) { + FES.userError('type error', + 'The first argument to a transform function (translate, rotate, scale, ' + + 'skewX, skewY) must be a transform created with transform2D() or transform3D().'); + } + const values = t.dimension === 4 ? values3D : values2D; + const m = build.matrixConstructorNode(strandsContext, t.dimension, values); + return t.mult(createStrandsNode(m.id, m.dimension, strandsContext)); + }; + + const originalTranslate = fn.translate; + augmentFn(fn, p5, 'translate', function (...args) { + const t = args[0]; + if (!isStrandsTransform(t)) return originalTranslate.apply(this, args); + const [, x = 0, y = 0, z = 0] = args; + return transformStep(t, + [1, 0, 0, 0, 1, 0, x, y, 1], + [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, x, y, z, 1]); + }); + + const originalScale = fn.scale; + augmentFn(fn, p5, 'scale', function (...args) { + const t = args[0]; + if (!isStrandsTransform(t)) return originalScale.apply(this, args); + const scales = args.slice(1); + const uniform = scales.length === 1; // scale(t, s) scales every axis by s + const x = scales[0] ?? 1; + const y = scales[1] ?? (uniform ? x : 1); + const z = scales[2] ?? (uniform ? x : 1); + return transformStep(t, + [x, 0, 0, 0, y, 0, 0, 0, 1], + [x, 0, 0, 0, 0, y, 0, 0, 0, 0, z, 0, 0, 0, 0, 1]); + }); + + const originalRotate = fn.rotate; + augmentFn(fn, p5, 'rotate', function (...args) { + const t = args[0]; + if (!isStrandsTransform(t)) return originalRotate.apply(this, args); + const angle = args[1]; // 2D: rotate in-plane; 3D: rotate about the Z axis + const c = this.cos(angle); + const s = this.sin(angle); + const ns = s.mult(-1); // -sin + return transformStep(t, + [c, s, 0, ns, c, 0, 0, 0, 1], + [c, s, 0, 0, ns, c, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]); + }); + + augmentFn(fn, p5, 'skewX', function (...args) { + if (!strandsContext.active) { + p5._friendlyError(`It looks like you've called skewX outside of a shader's modify() function.`); + return; + } + const [t, angle] = args; + const k = this.tan(angle); // x' = x + tan(angle) * y + return transformStep(t, + [1, 0, 0, k, 1, 0, 0, 0, 1], + [1, 0, 0, 0, k, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]); + }); + + augmentFn(fn, p5, 'skewY', function (...args) { + if (!strandsContext.active) { + p5._friendlyError(`It looks like you've called skewY outside of a shader's modify() function.`); + return; + } + const [t, angle] = args; + const k = this.tan(angle); // y' = y + tan(angle) * x + return transformStep(t, + [1, k, 0, 0, 1, 0, 0, 0, 1], + [1, k, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]); + }); + + augmentFn(fn, p5, 'rotateAxisAngle', function (...args) { + if (!strandsContext.active) { + p5._friendlyError(`It looks like you've called rotateAxisAngle outside of a shader's modify() function.`); + return; + } + const [t, axis, angle] = args; + if (!t?.isStrandsNode || t.typeInfo().baseType !== BaseType.MAT || t.dimension !== 4) { + FES.userError('type error', + 'rotateAxisAngle() only works on a 3D transform created with transform3D().'); + } + const n = this.normalize(axis); + const x = n.x, y = n.y, z = n.z; + const c = this.cos(angle); + const s = this.sin(angle); + const omc = p5.strandsNode(1).sub(c); // 1 - cos + + // Column-major Rodrigues' rotation matrix. + const b00 = x.mult(x).mult(omc).add(c); // x^2 * (1 - cos) + cos + const b01 = y.mult(x).mult(omc).add(z.mult(s)); // y * x * (1 - cos) + z * sin + const b02 = z.mult(x).mult(omc).sub(y.mult(s)); // z * x * (1 - cos) - y * sin + const b10 = x.mult(y).mult(omc).sub(z.mult(s)); // x * y * (1 - cos) - z * sin + const b11 = y.mult(y).mult(omc).add(c); // y^2 * (1 - cos) + cos + const b12 = z.mult(y).mult(omc).add(x.mult(s)); // z * y * (1 - cos) + x * sin + const b20 = x.mult(z).mult(omc).add(y.mult(s)); // x * z * (1 - cos) + y * sin + const b21 = y.mult(z).mult(omc).sub(x.mult(s)); // y * z * (1 - cos) - x * sin + const b22 = z.mult(z).mult(omc).add(c); // z^2 * (1 - cos) + cos + + const m = build.matrixConstructorNode(strandsContext, 4, [ + b00, b01, b02, 0, + b10, b11, b12, 0, + b20, b21, b22, 0, + 0, 0, 0, 1, + ]); + return t.mult(createStrandsNode(m.id, m.dimension, strandsContext)); + }); + + const original3DRotations = { + rotateX: fn.rotateX, + rotateY: fn.rotateY, + rotateZ: fn.rotateZ, + }; + const registerAxisRotation = (name, valuesFor) => { + const original = original3DRotations[name]; + augmentFn(fn, p5, name, function (...args) { + const t = args[0]; + if (!isStrandsTransform(t)) { + return original ? original.apply(this, args) : undefined; + } + if (t.dimension !== 4) { + FES.userError('type error', `${name}() needs a 3D transform created with transform3D().`); + } + const angle = args[1]; + const c = this.cos(angle); + const s = this.sin(angle); + const ns = s.mult(-1); + const m = build.matrixConstructorNode(strandsContext, 4, valuesFor(c, s, ns)); + return t.mult(createStrandsNode(m.id, m.dimension, strandsContext)); + }); + }; + registerAxisRotation('rotateX', (c, s, ns) => + [1, 0, 0, 0, 0, c, s, 0, 0, ns, c, 0, 0, 0, 0, 1]); + registerAxisRotation('rotateY', (c, s, ns) => + [c, 0, ns, 0, 0, 1, 0, 0, s, 0, c, 0, 0, 0, 0, 1]); + registerAxisRotation('rotateZ', (c, s, ns) => + [c, s, 0, 0, ns, c, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]); + + augmentFn(fn, p5, 'transformPoint', function (...args) { + if (!strandsContext.active) { + p5._friendlyError(`It looks like you've called transformPoint outside of a shader's modify() function.`); + return; + } + const [t, point] = args; + if (!(t?.isStrandsNode && t.typeInfo().baseType === BaseType.MAT)) { + FES.userError('type error', + 'transformPoint(t, point): the first argument must be a transform from transform2D()/transform3D().'); + } + const p = p5.strandsNode(point); + if (t.dimension === 4) { + return t.mult(fn.vec4(p, 1)).xyz; + } + return t.mult(fn.vec3(p, 1)).xy; + }); + + augmentFn(fn, p5, 'transformNormal', function (...args) { + if (!strandsContext.active) { + p5._friendlyError(`It looks like you've called transformNormal outside of a shader's modify() function.`); + return; + } + const [t, normal] = args; + if (!(t?.isStrandsNode && t.typeInfo().baseType === BaseType.MAT)) { + FES.userError('type error', + 'transformNormal(t, normal): the first argument must be a transform from transform2D()/transform3D().'); + } + const n = p5.strandsNode(normal); + // Normals ignore translation, so we use the transform's upper-left block. + // Its inverse-transpose is the correct normal matrix, staying accurate even + // under non-uniform scale. + const linear = fn[`mat${t.dimension - 1}`](t); + return fn.normalize(fn.transpose(fn.inverse(linear)).mult(n)); + }); + // Storage buffer uniform function for compute shaders fn.uniformStorage = function(name, bufferOrSchema) { let schema = null; diff --git a/src/strands/strands_builtins.js b/src/strands/strands_builtins.js index d6080ea428..e7b6d93e43 100644 --- a/src/strands/strands_builtins.js +++ b/src/strands/strands_builtins.js @@ -105,6 +105,18 @@ const builtInGLSLFunctions = { ], reflect: [{ params: [GenType.FLOAT, GenType.FLOAT], returnType: GenType.FLOAT, isp5Function: false}], refract: [{ params: [GenType.FLOAT, GenType.FLOAT,DataType.float1], returnType: GenType.FLOAT, isp5Function: false}], + + ////////// Matrix ////////// + inverse: [ + { params: [DataType.mat2], returnType: DataType.mat2, isp5Function: false}, + { params: [DataType.mat3], returnType: DataType.mat3, isp5Function: false}, + { params: [DataType.mat4], returnType: DataType.mat4, isp5Function: false}, + ], + transpose: [ + { params: [DataType.mat2], returnType: DataType.mat2, isp5Function: false}, + { params: [DataType.mat3], returnType: DataType.mat3, isp5Function: false}, + { params: [DataType.mat4], returnType: DataType.mat4, isp5Function: false}, + ], } export const strandsBuiltinFunctions = {