1import path from "node:path";
2import type * as rolldown from "rolldown";
3import { RolldownMagicString } from "rolldown";
4
5export const InternalWorkerExportPlugin = (): rolldown.Plugin => {
6 const getChunks = (bundle: rolldown.OutputBundle) => {
7 const chunks = new Map<string, rolldown.OutputChunk>();
8 for (const [fileName, output] of Object.entries(bundle)) {
9 if (output.type === "chunk") {
10 chunks.set(fileName, output);
11 }
12 }
13 return chunks;
14 };
15
16 const getImportedChunks = (
17 chunk: rolldown.OutputChunk,
18 chunks: Map<string, rolldown.OutputChunk>,
19 ) => {
20 const modules: Array<rolldown.OutputChunk> = [];
21 const seen = new Set<string>();
22
23 const visit = (module: rolldown.OutputChunk) => {
24 for (const fileName of module.imports) {
25 const imported = chunks.get(fileName);
26 if (!imported || seen.has(fileName)) {
27 continue;
28 }
29 seen.add(fileName);
30 visit(imported);
31 modules.push(imported);
32 }
33 };
34
35 visit(chunk);
36 return modules;
37 };
38
39 const getRelativeModulePath = (importer: string, imported: string) => {
40 const relativePath = path.posix.relative(
41 path.posix.dirname(importer),
42 imported,
43 );
44 return relativePath.startsWith(".") ? relativePath : `./${relativePath}`;
45 };
46
47 return {
48 name: "distilled:worker-exports",
49 generateBundle(_options, bundle) {
50 const chunks = getChunks(bundle);
51
52 for (const chunk of chunks.values()) {
53 const magicString = new RolldownMagicString(chunk.code);
54 const code = chunk.code;
55
56 if (chunk.isEntry) {
57 const modules = getImportedChunks(chunk, chunks);
58 const imports: Array<string> = [];
59 const moduleEntries: Array<string> = [];
60 for (let index = 0; index < modules.length; index++) {
61 const module = modules[index];
62 const relativePath = getRelativeModulePath(
63 chunk.fileName,
64 module.fileName,
65 );
66 if (module.isEntry) {
67 imports.push(
68 `import { modules as worker${index}Modules } from ${JSON.stringify(relativePath)};`,
69 );
70 moduleEntries.push(` ...worker${index}Modules,`);
71 } else {
72 imports.push(
73 `import worker${index} from ${JSON.stringify(relativePath)};`,
74 );
75 moduleEntries.push(
76 ` ${JSON.stringify(module.fileName)}: worker${index},`,
77 );
78 }
79 }
80
81 magicString.update(
82 0,
83 code.length,
84 [
85 ...imports,
86 imports.length > 0 ? "" : undefined,
87 `export const main = ${JSON.stringify(chunk.fileName)};`,
88 "export const modules = {",
89 ` ${JSON.stringify(chunk.fileName)}: ${JSON.stringify(code)},`,
90 ...moduleEntries,
91 "};",
92 "",
93 ]
94 .filter((line) => line !== undefined)
95 .join("\n"),
96 );
97 } else {
98 magicString.update(
99 0,
100 code.length,
101 `export default ${JSON.stringify(code)};\n`,
102 );
103 }
104
105 chunk.code = magicString.toString();
106 }
107 },
108 };
109};