Create typegpu babel plugin

apps/typegpu-docs/astro.config.mjs

@@ -5,9 +5,9 @@ import sitemap from '@astrojs/sitemap';
 import starlight from '@astrojs/starlight';
 import tailwind from '@astrojs/tailwind';
 import { defineConfig } from 'astro/config';
-import typegpu from 'rollup-plugin-typegpu';
 import starlightBlog from 'starlight-blog';
 import starlightTypeDoc, { typeDocSidebarGroup } from 'starlight-typedoc';
+import typegpu from 'unplugin-typegpu/rollup';
 import importRawRedirectPlugin from './vite-import-raw-redirect-plugin.mjs';
 
 /**

apps/typegpu-docs/package.json

@@ -15,7 +15,7 @@
     "@astrojs/starlight": "^0.25.2",
     "@astrojs/starlight-tailwind": "^2.0.3",
     "@astrojs/tailwind": "^5.1.0",
-    "@babel/standalone": "^7.24.7",
+    "@babel/standalone": "^7.26.6",
     "@monaco-editor/react": "^4.6.0",
     "@radix-ui/react-select": "^2.1.1",
     "@radix-ui/react-slider": "^1.2.0",
@@ -34,7 +34,7 @@
     "react": "^18.3.1",
     "react-dom": "^18.3.1",
     "remeda": "^2.3.0",
-    "rollup-plugin-typegpu": "workspace:*",
+    "unplugin-typegpu": "workspace:*",
     "sharp": "^0.32.5",
     "starlight-blog": "^0.12.0",
     "starlight-typedoc": "^0.17.0",
@@ -48,7 +48,7 @@
     "zod": "^3.23.8"
   },
   "devDependencies": {
-    "@types/babel__standalone": "^7.1.7",
+    "@types/babel__standalone": "^7.1.9",
     "@types/babel__template": "^7.4.4",
     "@types/babel__traverse": "^7.20.6",
     "@webgpu/types": "^0.1.43",

apps/typegpu-docs/src/utils/examples/exampleRunner.ts

@@ -54,13 +54,19 @@ const staticToDynamicImports = {
       const wildCard = imports.wildCard;
       const nonWildCard = imports.nonWildCard;
 
-      path.replaceWith(
-        template.program.ast(
-          `
-            ${wildCard?.length ? `const ${wildCard[0][1]} = await _import('${moduleName}');` : ''}
-            ${nonWildCard?.length ? `const { ${nonWildCard.map((imp) => (imp[0] === imp[1] ? imp[0] : `${imp[0]}: ${imp[1]}`)).join(',')} } = await _import('${moduleName}');` : ''}
-          `,
-        ),
+      path.replaceWithMultiple(
+        [
+          wildCard?.length
+            ? [
+                template.statement`const ${wildCard[0][1]} = await _import('${moduleName}');`(),
+              ]
+            : [],
+          nonWildCard?.length
+            ? [
+                template.statement`const { ${nonWildCard.map((imp) => (imp[0] === imp[1] ? imp[0] : `${imp[0]}: ${imp[1]}`)).join(',')} } = await _import('${moduleName}');`(),
+              ]
+            : [],
+        ].flat(),
       );
     },
   } satisfies TraverseOptions,
@@ -76,24 +82,21 @@ const exportedOptionsToExampleControls = () => {
 
         if (declaration?.type === 'VariableDeclaration') {
           const init = declaration.declarations[0].init;
+
           if (init) {
-            path.replaceWith(
-              template.program.ast(
-                `import { addParameters } from '@typegpu/example-toolkit';
-                addParameters(${code.slice(init.start ?? 0, init.end ?? 0)});`,
-              ),
-            );
+            path.replaceWithMultiple([
+              template.statement`import { addParameters } from '@typegpu/example-toolkit'`(),
+              template.statement`addParameters(${code.slice(init.start ?? 0, init.end ?? 0)});`(),
+            ]);
           }
         }
 
         if (declaration?.type === 'FunctionDeclaration') {
           const body = declaration.body;
-          path.replaceWith(
-            template.program.ast(
-              `import { onCleanup } from '@typegpu/example-toolkit';
-              onCleanup(() => ${code.slice(body.start ?? 0, body.end ?? 0)});`,
-            ),
-          );
+          path.replaceWithMultiple([
+            template.statement`import { onCleanup } from '@typegpu/example-toolkit'`(),
+            template.statement`onCleanup(() => ${code.slice(body.start ?? 0, body.end ?? 0)});`(),
+          ]);
         }
       },
     } satisfies TraverseOptions,

packages/rollup-plugin/README.md

@@ -2,7 +2,7 @@
 
 # rollup-plugin-typegpu
 
-🚧 **Under Construction** 🚧 - [GitHub](https://github.com/software-mansion/TypeGPU/tree/main/packages/rollup-plugin)
+⚠️ **This package is deprecated, please use `unplugin-typegpu` instead.** ⚠️
 
 </div>
 

packages/rollup-plugin/package.json

@@ -50,13 +50,10 @@
     "prepare-package": "tgpu-dev-cli prepack"
   },
   "dependencies": {
-    "tinyest-for-wgsl": "workspace:~0.1.0-alpha.0",
-    "estree-walker": "^3.0.3",
-    "magic-string": "^0.30.11"
+    "unplugin-typegpu": "workspace:^0.0.0"
   },
   "devDependencies": {
     "@typegpu/tgpu-dev-cli": "workspace:*",
-    "acorn": "^8.12.1",
     "rollup": "~4.12.0",
     "tsup": "^8.0.2",
     "typescript": "^5.3.3"

packages/rollup-plugin/src/index.ts

@@ -1,187 +1,4 @@
-import type { AnyNode, CallExpression } from 'acorn';
-import { walk } from 'estree-walker';
-import MagicString from 'magic-string';
-import type { Plugin, SourceMap } from 'rollup';
-import { transpileFn } from 'tinyest-for-wgsl';
+import rollupPlugin from 'unplugin-typegpu/rollup';
+export { type TypegpuPluginOptions } from 'unplugin-typegpu';
 
-const typegpuImportRegex = /import.*from\s*['"]typegpu.*['"]/g;
-const typegpuDynamicImportRegex = /import\s*\(\s*['"]\s*typegpu.*['"]/g;
-const typegpuRequireRegex = /require\s*\(\s*['"]\s*typegpu.*['"]\s*\)/g;
-
-type Context = {
-  /**
-   * How the `tgpu` object is used in code. Since it can be aliased, we
-   * need to catch that and act accordingly.
-   */
-  tgpuAliases: Set<string>;
-};
-
-type TgslFunctionDef = {
-  varDecl: CallExpression;
-  implementation: AnyNode;
-};
-
-function embedJSON(jsValue: unknown) {
-  return JSON.stringify(jsValue)
-    .replace(/\u2028/g, '\\u2028')
-    .replace(/\u2029/g, '\\u2029');
-}
-
-function gatherTgpuAliases(ctx: Context, node: AnyNode) {
-  if (node.type === 'ImportDeclaration') {
-    if (node.source.value === 'typegpu') {
-      for (const spec of node.specifiers) {
-        if (
-          // The default export of 'typegpu' is the `tgpu` object.
-          spec.type === 'ImportDefaultSpecifier' ||
-          // Aliasing 'tgpu' while importing, e.g. import { tgpu as t } from 'typegpu';
-          (spec.type === 'ImportSpecifier' &&
-            spec.imported.type === 'Identifier' &&
-            spec.imported.name === 'tgpu')
-        ) {
-          ctx.tgpuAliases.add(spec.local.name);
-        } else if (spec.type === 'ImportNamespaceSpecifier') {
-          // Importing everything, e.g. import * as t from 'typegpu';
-          ctx.tgpuAliases.add(`${spec.local.name}.tgpu`);
-        }
-      }
-    }
-  }
-}
-
-/**
- * Checks if `node` is an alias for the 'tgpu' object, traditionally
- * available via `import tgpu from 'typegpu'`.
- */
-function isTgpu(ctx: Context, node: AnyNode): boolean {
-  let path = '';
-
-  let tail = node;
-  while (true) {
-    if (tail.type === 'MemberExpression') {
-      if (tail.property.type !== 'Identifier') {
-        // Not handling computed expressions.
-        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 ctx.tgpuAliases.has(path);
-}
-
-export interface TypegpuPluginOptions {
-  include?: 'all' | RegExp[];
-}
-
-export interface TypegpuPlugin {
-  name: 'rollup-plugin-typegpu';
-  transform(
-    code: string,
-    id: string,
-  ): { code: string; map: SourceMap } | undefined;
-}
-
-export default function typegpu(options?: TypegpuPluginOptions): TypegpuPlugin {
-  return {
-    name: 'rollup-plugin-typegpu' as const,
-    transform(code, id) {
-      if (!options?.include) {
-        if (
-          !typegpuImportRegex.test(code) &&
-          !typegpuRequireRegex.test(code) &&
-          !typegpuDynamicImportRegex.test(code)
-        ) {
-          // No imports to `typegpu` or its sub modules, exiting early.
-          return;
-        }
-      } else if (
-        options.include !== 'all' &&
-        !options.include.some((pattern) => pattern.test(id))
-      ) {
-        return;
-      }
-
-      const ctx: Context = {
-        tgpuAliases: new Set(['tgpu']),
-      };
-
-      const ast = this.parse(code, {
-        allowReturnOutsideFunction: true,
-      });
-
-      const tgslFunctionDefs: TgslFunctionDef[] = [];
-
-      walk(ast, {
-        enter(_node, _parent, prop, index) {
-          const node = _node as AnyNode;
-
-          gatherTgpuAliases(ctx, node);
-
-          if (node.type === 'CallExpression') {
-            if (
-              node.callee.type === 'MemberExpression' &&
-              node.arguments.length === 1 &&
-              node.callee.property.type === 'Identifier' &&
-              ((node.callee.property.name === 'procedure' &&
-                isTgpu(ctx, node.callee.object)) ||
-                // Assuming that every call to `.does` is related to TypeGPU
-                // because shells can be created separately from calls to `tgpu`,
-                // making it hard to detect.
-                node.callee.property.name === 'does')
-            ) {
-              const implementation = node.arguments[0];
-
-              if (
-                implementation &&
-                !(implementation.type === 'TemplateLiteral') &&
-                !(implementation.type === 'Literal')
-              ) {
-                tgslFunctionDefs.push({
-                  varDecl: node,
-                  implementation,
-                });
-              }
-            }
-          }
-        },
-      });
-
-      const magicString = new MagicString(code);
-
-      for (const expr of tgslFunctionDefs) {
-        const { argNames, body, externalNames } = transpileFn(
-          expr.implementation,
-        );
-
-        // Wrap the implementation in a call to `tgpu.__assignAst` to associate the AST with the implementation.
-        magicString.appendLeft(expr.implementation.start, 'tgpu.__assignAst(');
-        magicString.appendRight(
-          expr.implementation.end,
-          `, ${embedJSON({ argNames, body, externalNames })}`,
-        );
-
-        if (externalNames.length > 0) {
-          magicString.appendRight(
-            expr.implementation.end,
-            `, {${externalNames.join(', ')}})`,
-          );
-        } else {
-          magicString.appendRight(expr.implementation.end, ', undefined)');
-        }
-      }
-
-      return {
-        code: magicString.toString(),
-        map: magicString.generateMap(),
-      };
-    },
-  } satisfies Plugin;
-}
+export default rollupPlugin;

packages/tinyest-for-wgsl/package.json

@@ -58,6 +58,7 @@
     "tsup": "^8.0.2",
     "typescript": "^5.3.3",
     "acorn": "^8.12.1",
+    "@babel/types": "7.26.5",
     "@typegpu/tgpu-dev-cli": "workspace:*"
   },
   "packageManager": "pnpm@8.15.8+sha256.691fe176eea9a8a80df20e4976f3dfb44a04841ceb885638fe2a26174f81e65e",