unplugin-typegpu
Version:
Build plugins for TypeGPU, enabling seamless JavaScript -> WGSL transpilation and improved debugging.
279 lines (277 loc) • 11.4 kB
JavaScript
import { transpileFn } from "tinyest-for-wgsl";
import * as t from "@babel/types";
import MagicString from "magic-string";
//#region src/core/common.ts
/**
* Each breaking change to the metadata format requires a bump to this number.
* It's used at runtime by `typegpu` to determine how to interpret a function's metadata.
*/
const METADATA_FORMAT_VERSION = 1;
function initPluginState(state, methods) {
var _state$opts$autoNamin;
state.tgpuAliases = new Set(state.opts.forceTgpuAlias ? [state.opts.forceTgpuAlias] : []);
state.autoNamingEnabled = (_state$opts$autoNamin = state.opts.autoNamingEnabled) !== null && _state$opts$autoNamin !== void 0 ? _state$opts$autoNamin : true;
state.inUseGpuScope = false;
state.alreadyTransformed = /* @__PURE__ */ new WeakSet();
Object.assign(state, methods);
}
/** Regular expressions used for early pruning (to avoid unnecessary parsing, which is expensive) */
const earlyPruneRegex = [/["']use gpu["']/, /t(ype)?gpu/];
const defaultOptions = {
include: /\.m?[jt]sx?(?:\?.*)?$/,
autoNamingEnabled: true,
earlyPruning: true
};
/**
* Returns the block scope of a function declaration, if one exists.
* Used to hoist a function declaration to the top of the scope it's visible in.
*/
function getBlockScope(path) {
if (!path.node.id) return;
const binding = path.scope.getBinding(path.node.id.name);
if (!binding) return;
let scopePath = binding.scope.path;
if (t.isNode(scopePath.node.body)) scopePath = scopePath.get("body");
if (t.isBlockStatement(scopePath.node) || t.isProgram(scopePath.node)) return scopePath;
}
/**
* Checks if `node` is an alias for the 'tgpu' object, traditionally
* available via `import tgpu from 'typegpu'`.
*/
function isTgpu(state, node) {
let path = "";
let tail = node;
while (true) if (tail.type === "MemberExpression") {
if (tail.property.type === "StringLiteral" && tail.property.value === "~unstable") {
tail = tail.object;
continue;
}
if (tail.property.type !== "Identifier") break;
path = path ? `${tail.property.name}.${path}` : tail.property.name;
tail = tail.object;
} else if (tail.type === "Identifier") {
path = path ? `${tail.name}.${path}` : tail.name;
break;
} else break;
return state.tgpuAliases.has(path);
}
function gatherTgpuAliases(state, node) {
if (node.source.value !== "typegpu") return;
for (const spec of node.specifiers) if (spec.type === "ImportDefaultSpecifier" || spec.type === "ImportSpecifier" && spec.imported.type === "Identifier" && spec.imported.name === "tgpu") state.tgpuAliases.add(spec.local.name);
else if (spec.type === "ImportNamespaceSpecifier") state.tgpuAliases.add(`${spec.local.name}.tgpu`);
}
const fnShellFunctionNames = [
"fn",
"vertexFn",
"fragmentFn",
"computeFn"
];
function isShellImplementationCall(node, state) {
return node.callee.type === "CallExpression" && node.callee.callee.type === "MemberExpression" && node.callee.callee.property.type === "Identifier" && fnShellFunctionNames.includes(node.callee.callee.property.name) && node.arguments.length === 1 && isTgpu(state, node.callee.callee.object);
}
/**
* Extracts a name and expression from nodes that contain a label and an expression,
* such as VariableDeclarator or PropertyDefinition.
* Returns a tuple of [name, expression] if found, otherwise undefined.
*
* @example
* extractLabelledExpression(node`let name = tgpu.bindGroupLayout({});`)
* // ["name", node`tgpu.bindGroupLayout({})`]
*/
function extractLabelledExpression(path) {
if (path.node.type === "VariableDeclarator" && path.node.id.type === "Identifier" && path.node.init) return [path.node.id.name, path.get("init")];
else if (path.node.type === "AssignmentExpression") {
const maybeName = tryFindIdentifier(path.node.left);
if (maybeName) return [maybeName, path.get("right")];
} else if (path.node.type === "ObjectProperty" && path.node.key.type === "Identifier") return [path.node.key.name, path.get("value")];
else if (path.node.type === "ClassProperty" && path.node.value && path.node.key.type === "Identifier") return [path.node.key.name, path.get("value")];
}
function getFunctionName(path) {
var _extractLabelledExpre, _path$node$id;
const maybeName = path.parentPath ? (_extractLabelledExpre = extractLabelledExpression(path.parentPath)) === null || _extractLabelledExpre === void 0 ? void 0 : _extractLabelledExpre[0] : void 0;
return maybeName !== null && maybeName !== void 0 ? maybeName : path.node.type === "FunctionDeclaration" || path.node.type === "FunctionExpression" ? (_path$node$id = path.node.id) === null || _path$node$id === void 0 ? void 0 : _path$node$id.name : void 0;
}
const resourceConstructors = [
"bindGroupLayout",
"vertexLayout",
"privateVar",
"workgroupVar",
"const",
"slot",
"accessor",
"mutableAccessor",
"comptime",
...fnShellFunctionNames,
"createBuffer",
"createMutable",
"createReadonly",
"createUniform",
"createQuerySet",
"createPipeline",
"createComputePipeline",
"createGuardedComputePipeline",
"createRenderPipeline",
"createTexture",
"createSampler",
"createComparisonSampler",
"struct",
"unstruct",
"createView"
];
/**
* Checks if `node` should be wrapped in an autoname function.
* Since it is mostly for debugging and clean WGSL generation,
* some false positives and false negatives are admissible.
*/
function containsResourceConstructorCall(node, state) {
if (node.type === "CallExpression") {
if (isShellImplementationCall(node, state)) return true;
if (node.callee.type === "Identifier" && resourceConstructors.includes(node.callee.name)) return true;
if (node.callee.type === "MemberExpression") {
if (node.callee.property.type === "Identifier") {
if (resourceConstructors.includes(node.callee.property.name)) return true;
if (node.callee.property.name === "$name") return false;
}
return containsResourceConstructorCall(node.callee.object, state);
}
}
if (node.type === "TaggedTemplateExpression") return containsResourceConstructorCall(node.tag, state);
return false;
}
/**
* Tries to find an identifier in a node.
*
* @example
* // syntax is simplified, imagine the arguments are appropriate nodes instead
* tryFindIdentifier('myBuffer'); // 'myBuffer'
* tryFindIdentifier('buffers.myBuffer'); // 'myBuffer'
* tryFindIdentifier('this.myBuffer'); // 'myBuffer'
* tryFindIdentifier('[a, b]'); // undefined
*/
function tryFindIdentifier(node) {
if (node.type === "Identifier") return node.name;
if (node.type === "PrivateName") return tryFindIdentifier(node.id);
if (node.type === "MemberExpression") return tryFindIdentifier(node.property);
}
/**
* Checks if `node` contains a label and a tgpu expression that could be named.
* If so, it calls the provided callback. Nodes selected for naming include:
*
* `let name = tgpu.bindGroupLayout({});` (VariableDeclarator)
*
* `name = tgpu.bindGroupLayout({});` (AssignmentExpression)
*
* `property: tgpu.bindGroupLayout({})` (Property/ObjectProperty)
*
* Since it is mostly for debugging and clean WGSL generation,
* some false positives and false negatives are admissible.
*
* This function is NOT used for auto-naming shell-less functions.
* Those are handled separately.
*
* @privateRemarks
* When adding new checks, you need to call this method in the corresponding node in Babel.
*/
function performExpressionNaming(state, path, namingCallback) {
if (!state.autoNamingEnabled) return;
const labelledExpression = extractLabelledExpression(path);
if (labelledExpression) {
const [label, expressionPath] = labelledExpression;
if (containsResourceConstructorCall(expressionPath.node, state)) namingCallback(expressionPath, label);
}
}
const operators = {
"+": "__tsover_add",
"-": "__tsover_sub",
"*": "__tsover_mul",
"/": "__tsover_div",
"%": "__tsover_mod",
"+=": "__tsover_add",
"-=": "__tsover_sub",
"*=": "__tsover_mul",
"/=": "__tsover_div",
"%=": "__tsover_mod"
};
function containsUseGpuDirective(node) {
var _node$body$directives, _node$body;
return ("directives" in node.body ? (_node$body$directives = (_node$body = node.body) === null || _node$body === void 0 ? void 0 : _node$body.directives) !== null && _node$body$directives !== void 0 ? _node$body$directives : [] : []).map((directive) => directive.value.value).includes("use gpu");
}
const fnNodeToTranspiledMap = /* @__PURE__ */ new WeakMap();
function functionOnExit(path, state) {
const node = path.node;
if (!containsUseGpuDirective(node)) return;
state.inUseGpuScope = false;
if (state.alreadyTransformed.has(node)) return;
const ast = fnNodeToTranspiledMap.get(path.node);
const maybeName = getFunctionName(path);
if (!ast) throw new Error(`No metadata found for function ${maybeName !== null && maybeName !== void 0 ? maybeName : "<unnamed>"}`);
state.assignMetadata(path, maybeName, ast);
state.alreadyTransformed.add(node);
path.skip();
}
const functionVisitor = {
ImportDeclaration(path, state) {
gatherTgpuAliases(state, path.node);
},
VariableDeclarator(path, state) {
performExpressionNaming(state, path, (pathToName, name) => {
state.wrapInAutoName(pathToName, name);
});
},
ObjectProperty(path, state) {
performExpressionNaming(state, path, (pathToName, name) => state.wrapInAutoName(pathToName, name));
},
ClassProperty(path, state) {
performExpressionNaming(state, path, (pathToName, name) => state.wrapInAutoName(pathToName, name));
},
AssignmentExpression: { exit(path, state) {
const runtimeFn = operators[path.node.operator];
if (state.inUseGpuScope && runtimeFn) state.replaceWithAssignmentOverload(path, runtimeFn);
performExpressionNaming(state, path, (pathToName, name) => state.wrapInAutoName(pathToName, name));
path.skip();
} },
BinaryExpression: { exit(path, state) {
const runtimeFn = operators[path.node.operator];
if (state.inUseGpuScope && runtimeFn) state.replaceWithBinaryOverload(path, runtimeFn);
path.skip();
} },
ArrowFunctionExpression: {
enter(path, state) {
if (containsUseGpuDirective(path.node)) {
fnNodeToTranspiledMap.set(path.node, transpileFn(path.node));
if (state.inUseGpuScope) throw new Error(`Nesting 'use gpu' functions is not allowed`);
state.inUseGpuScope = true;
}
},
exit: functionOnExit
},
FunctionExpression: {
enter(path, state) {
if (containsUseGpuDirective(path.node)) {
fnNodeToTranspiledMap.set(path.node, transpileFn(path.node));
if (state.inUseGpuScope) throw new Error(`Nesting 'use gpu' functions is not allowed`);
state.inUseGpuScope = true;
}
},
exit: functionOnExit
},
FunctionDeclaration: {
enter(path, state) {
if (containsUseGpuDirective(path.node)) {
fnNodeToTranspiledMap.set(path.node, transpileFn(path.node));
if (state.inUseGpuScope) throw new Error(`Nesting 'use gpu' functions is not allowed`);
state.inUseGpuScope = true;
}
},
exit: functionOnExit
},
CallExpression: { exit(path, state) {
const node = path.node;
if (isShellImplementationCall(node, state)) {
const implementation = node.arguments[0];
if (implementation && (implementation.type === "FunctionExpression" || implementation.type === "ArrowFunctionExpression") && !containsUseGpuDirective(implementation)) state.assignMetadata(path.get("arguments.0"), getFunctionName(path.get("arguments.0")), transpileFn(implementation));
}
} }
};
//#endregion
export { getBlockScope as a, functionVisitor as i, defaultOptions as n, initPluginState as o, earlyPruneRegex as r, METADATA_FORMAT_VERSION as t };