|
1 | | -import type { AnyNode, CallExpression } from 'acorn'; |
2 | | -import { walk } from 'estree-walker'; |
3 | | -import MagicString from 'magic-string'; |
4 | | -import type { Plugin, SourceMap } from 'rollup'; |
5 | | -import { transpileFn } from 'tinyest-for-wgsl'; |
| 1 | +import rollupPlugin from 'unplugin-typegpu/rollup'; |
| 2 | +export { type TypegpuPluginOptions } from 'unplugin-typegpu'; |
6 | 3 |
|
7 | | -const typegpuImportRegex = /import.*from\s*['"]typegpu.*['"]/g; |
8 | | -const typegpuDynamicImportRegex = /import\s*\(\s*['"]\s*typegpu.*['"]/g; |
9 | | -const typegpuRequireRegex = /require\s*\(\s*['"]\s*typegpu.*['"]\s*\)/g; |
10 | | - |
11 | | -type Context = { |
12 | | - /** |
13 | | - * How the `tgpu` object is used in code. Since it can be aliased, we |
14 | | - * need to catch that and act accordingly. |
15 | | - */ |
16 | | - tgpuAliases: Set<string>; |
17 | | -}; |
18 | | - |
19 | | -type TgslFunctionDef = { |
20 | | - varDecl: CallExpression; |
21 | | - implementation: AnyNode; |
22 | | -}; |
23 | | - |
24 | | -function embedJSON(jsValue: unknown) { |
25 | | - return JSON.stringify(jsValue) |
26 | | - .replace(/\u2028/g, '\\u2028') |
27 | | - .replace(/\u2029/g, '\\u2029'); |
28 | | -} |
29 | | - |
30 | | -function gatherTgpuAliases(ctx: Context, node: AnyNode) { |
31 | | - if (node.type === 'ImportDeclaration') { |
32 | | - if (node.source.value === 'typegpu') { |
33 | | - for (const spec of node.specifiers) { |
34 | | - if ( |
35 | | - // The default export of 'typegpu' is the `tgpu` object. |
36 | | - spec.type === 'ImportDefaultSpecifier' || |
37 | | - // Aliasing 'tgpu' while importing, e.g. import { tgpu as t } from 'typegpu'; |
38 | | - (spec.type === 'ImportSpecifier' && |
39 | | - spec.imported.type === 'Identifier' && |
40 | | - spec.imported.name === 'tgpu') |
41 | | - ) { |
42 | | - ctx.tgpuAliases.add(spec.local.name); |
43 | | - } else if (spec.type === 'ImportNamespaceSpecifier') { |
44 | | - // Importing everything, e.g. import * as t from 'typegpu'; |
45 | | - ctx.tgpuAliases.add(`${spec.local.name}.tgpu`); |
46 | | - } |
47 | | - } |
48 | | - } |
49 | | - } |
50 | | -} |
51 | | - |
52 | | -/** |
53 | | - * Checks if `node` is an alias for the 'tgpu' object, traditionally |
54 | | - * available via `import tgpu from 'typegpu'`. |
55 | | - */ |
56 | | -function isTgpu(ctx: Context, node: AnyNode): boolean { |
57 | | - let path = ''; |
58 | | - |
59 | | - let tail = node; |
60 | | - while (true) { |
61 | | - if (tail.type === 'MemberExpression') { |
62 | | - if (tail.property.type !== 'Identifier') { |
63 | | - // Not handling computed expressions. |
64 | | - break; |
65 | | - } |
66 | | - |
67 | | - path = path ? `${tail.property.name}.${path}` : tail.property.name; |
68 | | - tail = tail.object; |
69 | | - } else if (tail.type === 'Identifier') { |
70 | | - path = path ? `${tail.name}.${path}` : tail.name; |
71 | | - break; |
72 | | - } else { |
73 | | - break; |
74 | | - } |
75 | | - } |
76 | | - |
77 | | - return ctx.tgpuAliases.has(path); |
78 | | -} |
79 | | - |
80 | | -export interface TypegpuPluginOptions { |
81 | | - include?: 'all' | RegExp[]; |
82 | | -} |
83 | | - |
84 | | -export interface TypegpuPlugin { |
85 | | - name: 'rollup-plugin-typegpu'; |
86 | | - transform( |
87 | | - code: string, |
88 | | - id: string, |
89 | | - ): { code: string; map: SourceMap } | undefined; |
90 | | -} |
91 | | - |
92 | | -export default function typegpu(options?: TypegpuPluginOptions): TypegpuPlugin { |
93 | | - return { |
94 | | - name: 'rollup-plugin-typegpu' as const, |
95 | | - transform(code, id) { |
96 | | - if (!options?.include) { |
97 | | - if ( |
98 | | - !typegpuImportRegex.test(code) && |
99 | | - !typegpuRequireRegex.test(code) && |
100 | | - !typegpuDynamicImportRegex.test(code) |
101 | | - ) { |
102 | | - // No imports to `typegpu` or its sub modules, exiting early. |
103 | | - return; |
104 | | - } |
105 | | - } else if ( |
106 | | - options.include !== 'all' && |
107 | | - !options.include.some((pattern) => pattern.test(id)) |
108 | | - ) { |
109 | | - return; |
110 | | - } |
111 | | - |
112 | | - const ctx: Context = { |
113 | | - tgpuAliases: new Set(['tgpu']), |
114 | | - }; |
115 | | - |
116 | | - const ast = this.parse(code, { |
117 | | - allowReturnOutsideFunction: true, |
118 | | - }); |
119 | | - |
120 | | - const tgslFunctionDefs: TgslFunctionDef[] = []; |
121 | | - |
122 | | - walk(ast, { |
123 | | - enter(_node, _parent, prop, index) { |
124 | | - const node = _node as AnyNode; |
125 | | - |
126 | | - gatherTgpuAliases(ctx, node); |
127 | | - |
128 | | - if (node.type === 'CallExpression') { |
129 | | - if ( |
130 | | - node.callee.type === 'MemberExpression' && |
131 | | - node.arguments.length === 1 && |
132 | | - node.callee.property.type === 'Identifier' && |
133 | | - ((node.callee.property.name === 'procedure' && |
134 | | - isTgpu(ctx, node.callee.object)) || |
135 | | - // Assuming that every call to `.does` is related to TypeGPU |
136 | | - // because shells can be created separately from calls to `tgpu`, |
137 | | - // making it hard to detect. |
138 | | - node.callee.property.name === 'does') |
139 | | - ) { |
140 | | - const implementation = node.arguments[0]; |
141 | | - |
142 | | - if ( |
143 | | - implementation && |
144 | | - !(implementation.type === 'TemplateLiteral') && |
145 | | - !(implementation.type === 'Literal') |
146 | | - ) { |
147 | | - tgslFunctionDefs.push({ |
148 | | - varDecl: node, |
149 | | - implementation, |
150 | | - }); |
151 | | - } |
152 | | - } |
153 | | - } |
154 | | - }, |
155 | | - }); |
156 | | - |
157 | | - const magicString = new MagicString(code); |
158 | | - |
159 | | - for (const expr of tgslFunctionDefs) { |
160 | | - const { argNames, body, externalNames } = transpileFn( |
161 | | - expr.implementation, |
162 | | - ); |
163 | | - |
164 | | - // Wrap the implementation in a call to `tgpu.__assignAst` to associate the AST with the implementation. |
165 | | - magicString.appendLeft(expr.implementation.start, 'tgpu.__assignAst('); |
166 | | - magicString.appendRight( |
167 | | - expr.implementation.end, |
168 | | - `, ${embedJSON({ argNames, body, externalNames })}`, |
169 | | - ); |
170 | | - |
171 | | - if (externalNames.length > 0) { |
172 | | - magicString.appendRight( |
173 | | - expr.implementation.end, |
174 | | - `, {${externalNames.join(', ')}})`, |
175 | | - ); |
176 | | - } else { |
177 | | - magicString.appendRight(expr.implementation.end, ', undefined)'); |
178 | | - } |
179 | | - } |
180 | | - |
181 | | - return { |
182 | | - code: magicString.toString(), |
183 | | - map: magicString.generateMap(), |
184 | | - }; |
185 | | - }, |
186 | | - } satisfies Plugin; |
187 | | -} |
| 4 | +export default rollupPlugin; |
0 commit comments