tinyest-for-wgsl
Version:
Transforms JavaScript into its 'tinyest' form, to be used in generating equivalent (or close to) WGSL code.
313 lines (311 loc) • 9.93 kB
JavaScript
import * as tinyest from "tinyest";
import { FuncParameterType } from "tinyest";
//#region src/parsers.ts
const { NodeTypeCatalog: NODE } = tinyest;
function isDeclared(ctx, name) {
return ctx.stack.some((scope) => scope.declaredNames.includes(name));
}
const tsFallthrough = (ctx, node) => {
return transpile(ctx, node.expression);
};
const Transpilers = {
Program(ctx, node) {
const body = node.body[0];
if (!body) throw new Error("tgpu.fn was not implemented correctly.");
return transpile(ctx, body);
},
ExpressionStatement: (ctx, node) => transpile(ctx, node.expression),
ArrowFunctionExpression: () => {
throw new Error("Arrow functions are not supported inside TGSL.");
},
BlockStatement(ctx, node) {
ctx.stack.push({ declaredNames: [] });
const result = [NODE.block, node.body.map((statement) => transpile(ctx, statement))];
ctx.stack.pop();
return result;
},
ReturnStatement: (ctx, node) => node.argument ? [NODE.return, transpile(ctx, node.argument)] : [NODE.return],
Identifier(ctx, node) {
if (ctx.ignoreExternalDepth === 0 && !isDeclared(ctx, node.name)) ctx.externalNames.add(node.name);
return node.name;
},
ThisExpression(ctx) {
ctx.externalNames.add("this");
return "this";
},
BinaryExpression(ctx, node) {
const left = transpile(ctx, node.left);
const right = transpile(ctx, node.right);
return [
NODE.binaryExpr,
left,
node.operator,
right
];
},
LogicalExpression(ctx, node) {
const left = transpile(ctx, node.left);
const right = transpile(ctx, node.right);
return [
NODE.logicalExpr,
left,
node.operator,
right
];
},
AssignmentExpression(ctx, node) {
const left = transpile(ctx, node.left);
const right = transpile(ctx, node.right);
return [
NODE.assignmentExpr,
left,
node.operator,
right
];
},
UnaryExpression(ctx, node) {
const wgslOp = node.operator;
const argument = transpile(ctx, node.argument);
return [
NODE.unaryExpr,
wgslOp,
argument
];
},
MemberExpression(ctx, node) {
const object = transpile(ctx, node.object);
if (node.computed) {
const property$1 = transpile(ctx, node.property);
return [
NODE.indexAccess,
object,
property$1
];
}
ctx.ignoreExternalDepth++;
const property = transpile(ctx, node.property);
ctx.ignoreExternalDepth--;
if (typeof property !== "string") throw new Error("Expected identifier as property access key.");
return [
NODE.memberAccess,
object,
property
];
},
UpdateExpression(ctx, node) {
const operator = node.operator;
const argument = transpile(ctx, node.argument);
if (node.prefix) throw new Error("Prefix update expressions are not supported in WGSL.");
return [
NODE.postUpdate,
operator,
argument
];
},
ConditionalExpression(ctx, node) {
const test = transpile(ctx, node.test);
const consequent = transpile(ctx, node.consequent);
const alternative = transpile(ctx, node.alternate);
return [
NODE.conditionalExpr,
test,
consequent,
alternative
];
},
Literal(ctx, node) {
if (typeof node.value === "boolean") return node.value;
if (typeof node.value === "string") return [NODE.stringLiteral, node.value];
if (node.regex) throw new Error("Regular expression literals are not representable in WGSL.");
if (node.bigint) console.warn("BigInt literals are represented as numbers - loss of precision may occur.");
return [NODE.numericLiteral, String(Number(node.value))];
},
NumericLiteral(ctx, node) {
return [NODE.numericLiteral, String(node.value)];
},
BigIntLiteral(ctx, node) {
console.warn("BigInt literals are represented as numbers - loss of precision may occur.");
return [NODE.numericLiteral, String(Number.parseInt(node.value))];
},
BooleanLiteral(ctx, node) {
return node.value;
},
StringLiteral(ctx, node) {
return [NODE.stringLiteral, node.value];
},
CallExpression(ctx, node) {
const callee = transpile(ctx, node.callee);
const args = node.arguments.map((arg) => transpile(ctx, arg));
return [
NODE.call,
callee,
args
];
},
ArrayExpression: (ctx, node) => [NODE.arrayExpr, node.elements.map((elem) => {
if (!elem || elem.type === "SpreadElement") throw new Error("Spread elements are not supported in TGSL.");
return transpile(ctx, elem);
})],
VariableDeclaration(ctx, node) {
if (node.declarations.length !== 1 || !node.declarations[0]) throw new Error("Currently only one declaration in a statement is supported.");
const decl = node.declarations[0];
ctx.ignoreExternalDepth++;
const id = transpile(ctx, decl.id);
ctx.ignoreExternalDepth--;
if (typeof id !== "string") throw new Error("Invalid variable declaration, expected identifier.");
ctx.stack[ctx.stack.length - 1]?.declaredNames.push(id);
const init = decl.init ? transpile(ctx, decl.init) : void 0;
if (node.kind === "var") throw new Error("`var` declarations are not supported.");
if (node.kind === "const") return init !== void 0 ? [
NODE.const,
id,
init
] : [NODE.const, id];
return init !== void 0 ? [
NODE.let,
id,
init
] : [NODE.let, id];
},
IfStatement(ctx, node) {
const test = transpile(ctx, node.test);
const consequent = transpile(ctx, node.consequent);
const alternate = node.alternate ? transpile(ctx, node.alternate) : void 0;
return alternate ? [
NODE.if,
test,
consequent,
alternate
] : [
NODE.if,
test,
consequent
];
},
ObjectExpression(ctx, node) {
const properties = {};
for (const prop of node.properties) {
if (prop.type === "SpreadElement") throw new Error("Spread elements are not supported in TGSL.");
if (prop.key.type !== "Identifier" && prop.key.type !== "Literal") throw new Error("Only Identifier and Literal keys are supported as object keys.");
if (prop.type === "ObjectMethod") throw new Error("Object method elements are not supported in TGSL.");
ctx.ignoreExternalDepth++;
const key = prop.key.type === "Identifier" ? transpile(ctx, prop.key) : String(prop.key.value);
ctx.ignoreExternalDepth--;
properties[key] = transpile(ctx, prop.value);
}
return [NODE.objectExpr, properties];
},
ForStatement(ctx, node) {
const init = node.init ? transpile(ctx, node.init) : null;
const condition = node.test ? transpile(ctx, node.test) : null;
const update = node.update ? transpile(ctx, node.update) : null;
const body = transpile(ctx, node.body);
return [
NODE.for,
init,
condition,
update,
body
];
},
WhileStatement(ctx, node) {
const condition = transpile(ctx, node.test);
const body = transpile(ctx, node.body);
return [
NODE.while,
condition,
body
];
},
ForOfStatement(ctx, node) {
const loopVar = transpile(ctx, node.left);
const iterable = transpile(ctx, node.right);
const body = transpile(ctx, node.body);
return [
NODE.forOf,
loopVar,
iterable,
body
];
},
ContinueStatement() {
return [NODE.continue];
},
BreakStatement() {
return [NODE.break];
},
TSAsExpression: tsFallthrough,
TSSatisfiesExpression: tsFallthrough,
TSNonNullExpression: tsFallthrough
};
function transpile(ctx, node) {
const transpiler = Transpilers[node.type];
if (!transpiler) throw new Error(`Unsupported JS functionality: ${node.type}`);
return transpiler(ctx, node);
}
function extractFunctionParts(rootNode) {
let functionNode = null;
let unwrappedNode = rootNode;
while (true) if (unwrappedNode.type === "Program") {
const statement = unwrappedNode.body.filter((n) => n.type === "ExpressionStatement" || n.type === "FunctionDeclaration")[0];
if (!statement) break;
unwrappedNode = statement;
} else if (unwrappedNode.type === "ExpressionStatement") unwrappedNode = unwrappedNode.expression;
else if (unwrappedNode.type === "ArrowFunctionExpression") {
functionNode = unwrappedNode;
break;
} else if (unwrappedNode.type === "FunctionExpression") {
functionNode = unwrappedNode;
break;
} else if (unwrappedNode.type === "FunctionDeclaration") {
functionNode = unwrappedNode;
break;
} else break;
if (!functionNode) throw new Error(`tgpu.fn expected a single function to be passed as implementation ${JSON.stringify(unwrappedNode)}`);
if (functionNode.async) throw new Error("tgpu.fn cannot be async");
if (functionNode.generator) throw new Error("tgpu.fn cannot be a generator");
const unsupportedTypes = new Set(functionNode.params.flatMap((param) => param.type === "ObjectPattern" || param.type === "Identifier" ? [] : [param.type]));
if (unsupportedTypes.size > 0) throw new Error(`Unsupported function parameter type(s): ${[...unsupportedTypes].join(", ")}`);
return {
params: functionNode.params.map((param) => param.type === "ObjectPattern" ? {
type: FuncParameterType.destructuredObject,
props: param.properties.flatMap((prop) => (prop.type === "Property" || prop.type === "ObjectProperty") && prop.key.type === "Identifier" && prop.value.type === "Identifier" ? [{
name: prop.key.name,
alias: prop.value.name
}] : [])
} : {
type: FuncParameterType.identifier,
name: param.name
}),
body: functionNode.body
};
}
function transpileFn(rootNode) {
const { params, body } = extractFunctionParts(rootNode);
const ctx = {
externalNames: /* @__PURE__ */ new Set(),
ignoreExternalDepth: 0,
stack: [{ declaredNames: params.flatMap((param) => param.type === FuncParameterType.identifier ? param.name : param.props.map((prop) => prop.alias)) }]
};
const tinyestBody = transpile(ctx, body);
const externalNames = [...ctx.externalNames];
if (body.type === "BlockStatement") return {
params,
body: tinyestBody,
externalNames
};
return {
params,
body: [NODE.block, [[NODE.return, tinyestBody]]],
externalNames
};
}
function transpileNode(node) {
return transpile({
externalNames: /* @__PURE__ */ new Set(),
ignoreExternalDepth: 0,
stack: [{ declaredNames: [] }]
}, node);
}
//#endregion
export { transpileFn, transpileNode };