Create typegpu babel plugin

apps/typegpu-docs/astro.config.mjs

@@ -5,7 +5,7 @@ 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 typegpu from 'plugin-typegpu/rollup';
 import starlightBlog from 'starlight-blog';
 import starlightTypeDoc, { typeDocSidebarGroup } from 'starlight-typedoc';
 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",
@@ -32,7 +32,7 @@
     "react": "^18.3.1",
     "react-dom": "^18.3.1",
     "remeda": "^2.3.0",
-    "rollup-plugin-typegpu": "workspace:*",
+    "plugin-typegpu": "workspace:*",
     "sharp": "^0.32.5",
     "starlight-blog": "^0.12.0",
     "starlight-typedoc": "^0.17.0",
@@ -45,7 +45,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/plugin/README.md

@@ -0,0 +1,50 @@
+<div align="center">
+
+# plugin-typegpu
+
+🚧 **Under Construction** 🚧 - [GitHub](https://github.com/software-mansion/TypeGPU/tree/main/packages/plugin)
+
+</div>
+
+Build plugins for [TypeGPU](https://typegpu.com), transpiling JavaScript functions into WGSL at build-time.
+
+Currently supported:
+* babel
+```js
+// babel.config.js (React Native with Expo)
+
+module.exports = function (api) {
+  api.cache(true);
+  return {
+    presets: ['babel-preset-expo'],
+    plugins: ['plugin-typegpu/babel']
+  };
+};
+
+```
+
+* rollup
+
+```ts
+// vite.config.js
+
+import { defineConfig } from 'vite';
+import typegpu from 'plugin-typegpu/rollup';
+
+export default defineConfig({
+  plugins: [typegpu()],
+});
+
+```
+
+## Getting Started
+
+```sh
+npm install plugin-typegpu
+```
+
+## TypeGPU is created by Software Mansion
+
+[![swm](https://logo.swmansion.com/logo?color=white&variant=desktop&width=150&tag=typegpu-github 'Software Mansion')](https://swmansion.com)
+
+Since 2012 [Software Mansion](https://swmansion.com) is a software agency with experience in building web and mobile apps. We are Core React Native Contributors and experts in dealing with all kinds of React Native issues. We can help you build your next dream product – [Hire us](https://swmansion.com/contact/projects?utm_source=typegpu&utm_medium=readme).

packages/plugin/package.json

@@ -0,0 +1,77 @@
+{
+  "name": "plugin-typegpu",
+  "version": "0.0.0",
+  "description": "Build plugins for TypeGPU, transpiling JavaScript functions into WGSL at build-time.",
+  "keywords": [
+    "rollup-plugin",
+    "babel-plugin",
+    "vite-plugin",
+    "typegpu",
+    "gpgpu",
+    "webgpu",
+    "wgpu",
+    "wgsl",
+    "typescript",
+    "shaders"
+  ],
+  "license": "MIT",
+  "type": "module",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    "./package.json": "./package.json",
+    ".": {
+      "types": "./dist/index.d.ts",
+      "module": "./dist/index.js",
+      "import": "./dist/index.js",
+      "default": "./dist/index.cjs"
+    },
+    "./rollup": {
+      "types": "./dist/rollup.d.ts",
+      "module": "./dist/rollup.js",
+      "import": "./dist/rollup.js",
+      "default": "./dist/rollup.cjs"
+    },
+    "./babel": {
+      "types": "./dist/babel.d.ts",
+      "module": "./dist/babel.js",
+      "import": "./dist/babel.js",
+      "default": "./dist/babel.cjs"
+    }
+  },
+  "sideEffects": false,
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/software-mansion/TypeGPU.git"
+  },
+  "bugs": {
+    "url": "https://github.com/software-mansion/TypeGPU/issues"
+  },
+  "homepage": "https://typegpu.com",
+  "scripts": {
+    "dev:watch": "DEV=true tsup --watch",
+    "dev:build": "DEV=true tsup",
+    "build": "tsup",
+    "test:types": "pnpm tsc --p ./tsconfig.json --noEmit",
+    "publish": "echo \"Use pnpm prepare-package instead!\" && exit 1",
+    "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",
+    "@babel/standalone": "^7.26.6"
+  },
+  "devDependencies": {
+    "@babel/template": "^7.25.9",
+    "@babel/types": "^7.26.5",
+    "@typegpu/tgpu-dev-cli": "workspace:*",
+    "@types/babel__standalone": "^7.1.9",
+    "@types/babel__template": "^7.4.4",
+    "@types/babel__traverse": "^7.20.6",
+    "acorn": "^8.12.1",
+    "rollup": "~4.12.0",
+    "tsup": "^8.0.2",
+    "typescript": "^5.3.3"
+  }
+}

packages/plugin/src/babel.ts

@@ -0,0 +1,84 @@
+import * as Babel from '@babel/standalone';
+import type TemplateGenerator from '@babel/template';
+import type { TraverseOptions } from '@babel/traverse';
+import type * as t from '@babel/types';
+import { transpileFn } from 'tinyest-for-wgsl';
+import {
+  type Context,
+  type TypegpuPluginOptions,
+  embedJSON,
+  gatherTgpuAliases,
+  isDoesCall,
+  shouldSkipFile,
+} from './common';
+
+// NOTE: @babel/standalone does expose internal packages, as specified in the docs, but the
+// typing for @babel/standalone does not expose them.
+const template = (
+  Babel as unknown as { packages: { template: typeof TemplateGenerator } }
+).packages.template;
+const types = (Babel as unknown as { packages: { types: typeof t } }).packages
+  .types;
+
+function functionVisitor(ctx: Context): TraverseOptions {
+  return {
+    ImportDeclaration(path) {
+      gatherTgpuAliases(path.node, ctx);
+    },
+
+    CallExpression(path) {
+      const node = path.node;
+
+      if (isDoesCall(node, ctx)) {
+        const implementation = node.arguments[0];
+
+        if (
+          implementation &&
+          !(implementation.type === 'TemplateLiteral') &&
+          !(implementation.type === 'StringLiteral')
+        ) {
+          const { argNames, body, externalNames } = transpileFn(implementation);
+
+          path.replaceWith(
+            types.callExpression(node.callee, [
+              types.callExpression(template.expression('tgpu.__assignAst')(), [
+                implementation,
+                template.expression`${embedJSON({ argNames, body, externalNames })}`(),
+                externalNames.length > 0
+                  ? template.expression`{${externalNames.join(', ')}}`()
+                  : template.expression`undefined`(),
+              ]),
+            ]),
+          );
+
+          path.skip();
+        }
+      }
+    },
+  };
+}
+
+export default function () {
+  return {
+    visitor: {
+      Program(path, state) {
+        // @ts-ignore
+        const code: string | undefined = state.file?.code;
+        // @ts-ignore
+        const options: TypegpuPluginOptions | undefined = state.opts;
+        // @ts-ignore
+        const id: string | undefined = state.filename;
+
+        if (shouldSkipFile(options, id, code)) {
+          return;
+        }
+
+        const ctx: Context = {
+          tgpuAliases: new Set(['tgpu']),
+        };
+
+        path.traverse(functionVisitor(ctx));
+      },
+    } satisfies TraverseOptions,
+  };
+}