unplugin-typegpu
Version:
Build plugins for TypeGPU, enabling seamless JavaScript -> WGSL transpilation and improved debugging.
1 lines • 11 kB
Source Map (JSON)
{"version":3,"sources":["../src/babel.ts","../src/common.ts"],"names":["defaultOptions","embedJSON","jsValue","isTgpu","ctx","node","path","tail"],"mappings":"AAAA,whBAAuB,kCAIQ,kDACH,ICefA,CAAAA,CAAiB,CAC5B,OAAA,CAAS,CAAC,cAAc,CAC1B,CAAA,CAEO,SAASC,CAAAA,CAAUC,CAAAA,CAAkB,CAC1C,OAAO,IAAA,CAAK,SAAA,CAAUA,CAAO,CAAA,CAC1B,OAAA,CAAQ,SAAA,CAAW,SAAS,CAAA,CAC5B,OAAA,CAAQ,SAAA,CAAW,SAAS,CACjC,CAMA,SAASC,CAAAA,CAAOC,CAAAA,CAAcC,CAAAA,CAA2C,CACvE,IAAIC,CAAAA,CAAO,EAAA,CAEPC,CAAAA,CAAOF,CAAAA,CACX,GAAA,CAAA,CAAA,CAAA,CACE,EAAA,CAAIE,CAAAA,CAAK,IAAA,GAAS,kBAAA,CAAoB,CACpC,EAAA,CAAIA,CAAAA,CAAK,QAAA,CAAS,IAAA,GAAS,YAAA,CAEzB,KAAA,CAGFD,CAAAA,CAAOA,CAAAA,CAAO,CAAA,EAAA;ADcG,OAAA;AAC8B,SAAA;AACX,gBAAA;AA0BtB,GAAA","file":"/Users/iwo/Projects/wigsill/packages/unplugin-typegpu/dist/chunk-G6NIIU55.cjs","sourcesContent":["import * as Babel from '@babel/standalone';\nimport type TemplateGenerator from '@babel/template';\nimport type { TraverseOptions } from '@babel/traverse';\nimport type * as babel from '@babel/types';\nimport { FORMAT_VERSION } from 'tinyest';\nimport { transpileFn } from 'tinyest-for-wgsl';\nimport {\n type Context,\n embedJSON,\n gatherTgpuAliases,\n getErrorMessage,\n isShellImplementationCall,\n type KernelDirective,\n kernelDirectives,\n type Options,\n} from './common.ts';\nimport { createFilterForId } from './filter.ts';\n\n// NOTE: @babel/standalone does expose internal packages, as specified in the docs, but the\n// typing for @babel/standalone does not expose them.\nconst template = (\n Babel as unknown as { packages: { template: typeof TemplateGenerator } }\n).packages.template;\nconst types = (Babel as unknown as { packages: { types: typeof babel } })\n .packages.types;\n\nfunction getKernelDirective(\n node:\n | babel.FunctionDeclaration\n | babel.FunctionExpression\n | babel.ArrowFunctionExpression,\n): KernelDirective | undefined {\n const directives = (\n 'directives' in node.body ? (node.body?.directives ?? []) : []\n ).map((directive) => directive.value.value);\n\n for (const directive of kernelDirectives) {\n if (directives.includes(directive)) {\n return directive;\n }\n }\n}\n\nfunction i(identifier: string): babel.Identifier {\n return types.identifier(identifier);\n}\n\nfunction functionToTranspiled(\n node: babel.ArrowFunctionExpression | babel.FunctionExpression,\n directive: KernelDirective | undefined,\n name?: string | undefined,\n): babel.CallExpression | null {\n if (!directive) {\n return null;\n }\n\n const { params, body, externalNames } = transpileFn(node);\n\n const metadata = `{\n v: ${FORMAT_VERSION},\n ast: ${embedJSON({ params, body, externalNames })},\n externals: {${externalNames.join(', ')}},\n }`;\n\n const jsImpl = directive === 'kernel & js'\n ? node\n : types.arrowFunctionExpression(\n [],\n types.blockStatement(\n [types.throwStatement(\n types.newExpression(i('Error'), [\n types.stringLiteral(getErrorMessage(name)),\n ]),\n )],\n ),\n );\n\n return types.callExpression(\n types.arrowFunctionExpression(\n [i('$')],\n types.logicalExpression(\n '&&',\n types.callExpression(\n types.memberExpression(\n types.assignmentExpression(\n '??=',\n types.memberExpression(i('globalThis'), i('__TYPEGPU_META__')),\n types.newExpression(i('WeakMap'), []),\n ),\n i('set'),\n ),\n [\n types.assignmentExpression(\n '=',\n types.memberExpression(i('$'), i('f')),\n jsImpl,\n ),\n template.expression`${metadata}`(),\n ],\n ),\n types.memberExpression(i('$'), i('f')),\n ),\n ),\n [types.objectExpression([])],\n );\n}\n\nfunction functionVisitor(ctx: Context): TraverseOptions {\n return {\n ImportDeclaration(path) {\n gatherTgpuAliases(path.node, ctx);\n },\n\n ArrowFunctionExpression(path) {\n const transpiled = functionToTranspiled(\n path.node,\n getKernelDirective(path.node),\n path.parentPath.node.type === 'VariableDeclarator'\n ? path.parentPath.node.id.type === 'Identifier'\n ? path.parentPath.node.id.name\n : undefined\n : undefined,\n );\n if (transpiled) {\n path.replaceWith(transpiled);\n path.skip();\n }\n },\n\n FunctionExpression(path) {\n const transpiled = functionToTranspiled(\n path.node,\n getKernelDirective(path.node),\n path.node.id?.name\n ? path.node.id.name\n : path.parentPath.node.type === 'VariableDeclarator'\n ? path.parentPath.node.id.type === 'Identifier'\n ? path.parentPath.node.id.name\n : undefined\n : undefined,\n );\n if (transpiled) {\n path.replaceWith(transpiled);\n path.skip();\n }\n },\n\n FunctionDeclaration(path) {\n const node = path.node;\n const expression = types.functionExpression(\n node.id,\n node.params,\n node.body,\n );\n const transpiled = functionToTranspiled(\n expression,\n getKernelDirective(path.node),\n node.id?.name,\n );\n if (transpiled && node.id) {\n path.replaceWith(\n types.variableDeclaration('const', [\n types.variableDeclarator(node.id, transpiled),\n ]),\n );\n path.skip();\n }\n },\n\n CallExpression(path) {\n const node = path.node;\n\n if (isShellImplementationCall(node, ctx)) {\n const implementation = node.arguments[0];\n\n if (\n implementation &&\n (implementation.type === 'FunctionExpression' ||\n implementation.type === 'ArrowFunctionExpression')\n ) {\n const transpiled = functionToTranspiled(\n implementation,\n getKernelDirective(implementation) ?? 'kernel',\n ) as babel.CallExpression;\n\n path.replaceWith(\n types.callExpression(node.callee, [\n transpiled,\n ]),\n );\n\n path.skip();\n }\n }\n },\n };\n}\n\nexport default function () {\n return {\n visitor: {\n Program(path, state) {\n // biome-ignore lint/suspicious/noExplicitAny: <oh babel babel...>\n const code: string | undefined = (state as any).file?.code;\n // biome-ignore lint/suspicious/noExplicitAny: <oh babel babel...>\n const options: Options | undefined = (state as any).opts;\n // biome-ignore lint/suspicious/noExplicitAny: <oh babel babel...>\n const id: string | undefined = (state as any).filename;\n\n const filter = createFilterForId(options);\n if (id && filter && !filter?.(id)) {\n return;\n }\n\n const ctx: Context = {\n tgpuAliases: new Set<string>(\n options?.forceTgpuAlias ? [options.forceTgpuAlias] : [],\n ),\n fileId: id,\n };\n\n path.traverse(functionVisitor(ctx));\n },\n } satisfies TraverseOptions,\n };\n}\n","import type * as babel from '@babel/types';\nimport type * as acorn from 'acorn';\nimport type { FilterPattern } from 'unplugin';\n\nexport type Context = {\n /**\n * How the `tgpu` object is used in code. Since it can be aliased, we\n * need to catch that and act accordingly.\n */\n tgpuAliases: Set<string>;\n fileId?: string | undefined;\n};\n\nexport interface Options {\n include?: FilterPattern;\n exclude?: FilterPattern;\n enforce?: 'post' | 'pre' | undefined;\n forceTgpuAlias?: string;\n}\n\nexport const defaultOptions = {\n include: [/\\.m?[jt]sx?$/],\n};\n\nexport function embedJSON(jsValue: unknown) {\n return JSON.stringify(jsValue)\n .replace(/\\u2028/g, '\\\\u2028')\n .replace(/\\u2029/g, '\\\\u2029');\n}\n\n/**\n * Checks if `node` is an alias for the 'tgpu' object, traditionally\n * available via `import tgpu from 'typegpu'`.\n */\nfunction isTgpu(ctx: Context, node: babel.Node | acorn.AnyNode): boolean {\n let path = '';\n\n let tail = node;\n while (true) {\n if (tail.type === 'MemberExpression') {\n if (tail.property.type !== 'Identifier') {\n // Not handling computed expressions.\n break;\n }\n\n path = path ? `${tail.property.name}.${path}` : tail.property.name;\n tail = tail.object;\n } else if (tail.type === 'Identifier') {\n path = path ? `${tail.name}.${path}` : tail.name;\n break;\n } else {\n break;\n }\n }\n\n return ctx.tgpuAliases.has(path);\n}\n\nexport function gatherTgpuAliases(\n node: acorn.ImportDeclaration | babel.ImportDeclaration,\n ctx: Context,\n) {\n if (node.source.value === 'typegpu') {\n for (const spec of node.specifiers) {\n if (\n // The default export of 'typegpu' is the `tgpu` object.\n spec.type === 'ImportDefaultSpecifier' ||\n // Aliasing 'tgpu' while importing, e.g. import { tgpu as t } from 'typegpu';\n (spec.type === 'ImportSpecifier' &&\n spec.imported.type === 'Identifier' &&\n spec.imported.name === 'tgpu')\n ) {\n ctx.tgpuAliases.add(spec.local.name);\n } else if (spec.type === 'ImportNamespaceSpecifier') {\n // Importing everything, e.g. import * as t from 'typegpu';\n ctx.tgpuAliases.add(`${spec.local.name}.tgpu`);\n }\n }\n }\n}\n\nconst fnShellFunctionNames = ['fn', 'vertexFn', 'fragmentFn', 'computeFn'];\n\nexport function isShellImplementationCall(\n node: acorn.CallExpression | babel.CallExpression,\n ctx: Context,\n) {\n return (\n (node.callee.type === 'CallExpression' &&\n node.callee.callee.type === 'MemberExpression' &&\n node.callee.callee.property.type === 'Identifier' &&\n fnShellFunctionNames.includes(node.callee.callee.property.name) &&\n node.arguments.length === 1 &&\n (node.callee.callee.object.type === 'MemberExpression'\n ? isTgpu(ctx, node.callee.callee.object.object)\n : isTgpu(ctx, node.callee.callee.object))) || // TODO: remove along with the deprecated 'does' method\n (node.callee.type === 'MemberExpression' &&\n node.arguments.length === 1 &&\n node.callee.property.type === 'Identifier' &&\n // Assuming that every call to `.does` is related to TypeGPU\n // because shells can be created separately from calls to `tgpu`,\n // making it hard to detect.\n node.callee.property.name === 'does')\n );\n}\n\nexport const kernelDirectives = ['kernel', 'kernel & js'] as const;\nexport type KernelDirective = (typeof kernelDirectives)[number];\n\nexport function getErrorMessage(name: string | undefined) {\n return `The function \"${\n name ?? '<unnamed>'\n }\" is invokable only on the GPU. If you want to use it on the CPU, mark it with the \"kernel & js\" directive.`;\n}\n"]}