Perf: Use the AST and analyze the variables used to shim the ESM on d…

…emand

packages/techor/package.json

@@ -73,14 +73,14 @@
         "@techor/glob": "workspace:^",
         "@techor/log": "workspace:^",
         "@techor/npm": "workspace:^",
+        "acorn": "^8.11.3",
         "clsx": "^2.0.0",
         "escodegen": "^2.1.0",
         "execa": "^7.2.0",
         "explore-config": "workspace:^",
         "hrtime": "^0.5.0",
         "load-tsconfig": "^0.2.5",
         "lodash.isequal": "^4.5.0",
-        "magic-string": "^0.30.7",
         "pkg-types": "^1.0.1",
         "pretty-bytes": "^6.1.0",
         "pretty-hrtime": "^1.0.3",
@@ -94,6 +94,7 @@
         "yargs-parser": "^21.1.1"
     },
     "devDependencies": {
+        "@types/escodegen": "^0.0.10",
         "dedent": "^0.7.0"
     }
 }

packages/techor/src/plugins/esm-shim.ts

@@ -1,36 +1,90 @@
 import type { Plugin } from 'rollup'
-import MagicString from 'magic-string'
+import { parse, Options as AcornOptions, Statement, ModuleDeclaration } from 'acorn'
+import escodegen from 'escodegen'
 
-export const ESM_SHIM_CODE = `
-/* Techor ESM Shim */
-import __url_for_shim from 'url';
-import __path_for_shim from 'path';
-import __mod_for_shim from 'module';
-const __filename = __url_for_shim.fileURLToPath(import.meta.url);
-const __dirname = __path_for_shim.dirname(__filename);
-const require = __mod_for_shim.createRequire(import.meta.url);\n
-`
+const parseOptions: AcornOptions = { ecmaVersion: 'latest', sourceType: 'module' }
+
+const urlImportDeclaration = parse('import __url_for_shim from "url";', parseOptions).body[0]
+const pathImportDeclaration = parse('import __path_for_shim from "path";', parseOptions).body[0]
+const moduleImportDeclaration = parse('import __mod_for_shim from "module";', parseOptions).body[0]
+
+const filenameDeclaration = parse('const __filename = __url_for_shim.fileURLToPath(import.meta.url);', parseOptions).body[0]
+const dirnameDeclaration = parse('const __dirname = __path_for_shim.dirname(__filename);', parseOptions).body[0]
+const requireDeclaration = parse('const require = __mod_for_shim.createRequire(import.meta.url);', parseOptions).body[0]
 
 export default function esmShim(): Plugin {
     return {
         name: 'techor-esm-shim',
-        renderChunk(code, chunk, options) {
-            if (options.format === 'es' && !code.includes(ESM_SHIM_CODE) && /__filename|__dirname|require\(|require\.resolve\(/.test(code)) {
-                const str = new MagicString(code)
-                /* insert ESM_SHIM_CODE after the something like:
-                 * #!/usr/bin/env node
-                */
-                if (code.startsWith('#!')) {
-                    str.prependRight(code.indexOf('\n') + 1, ESM_SHIM_CODE)
-                } else {
-                    str.prepend(ESM_SHIM_CODE)
-                }
-                return {
-                    code: str.toString(),
-                    map: str.generateMap()
+        renderChunk: {
+            order: 'pre',
+            async handler(code, chunk, options) {
+                if (options.format === 'es') {
+                    const ast = this.parse(code)
+                    if (ast.type === 'Program' && ast.body) {
+                        const newDeclarations = new Set<Statement | ModuleDeclaration>()
+                        const filteredBody = ast.body.filter(Boolean)
+                        let dirnameDeclared = false
+                        let filenameDeclared = false
+                        let requireDeclared = false
+                        // top-level declarations
+                        for (const node of filteredBody) {
+                            if (node.type === 'VariableDeclaration') {
+                                const varName = node.declarations[0].id?.['name']
+                                switch (varName) {
+                                    case '__dirname':
+                                        dirnameDeclared = true
+                                        break
+                                    case '__filename':
+                                        filenameDeclared = true
+                                        break
+                                    case 'require':
+                                        requireDeclared = true
+                                        break
+                                }
+                            }
+                        }
+                        let dirnameUsed = false
+                        let filenameUsed = false
+                        let requireUsed = false
+                        // walk the ESTree to find whether the variables `__dirname` `__filename` `require` are used
+                        const walk = (node: any) => {
+                            if (node.type === 'Identifier') {
+                                const name = node.name
+                                if (name === '__dirname') dirnameUsed = true
+                                if (name === '__filename') filenameUsed = true
+                                if (name === 'require') requireUsed = true
+                            }
+                            for (const key in node) {
+                                if (node[key] && typeof node[key] === 'object') {
+                                    walk(node[key])
+                                }
+                            }
+                        }
+                        for (const node of filteredBody) {
+                            walk(node)
+                        }
+                        // import
+                        if (!filenameDeclared && filenameUsed || !dirnameDeclared && dirnameUsed) newDeclarations.add(urlImportDeclaration)
+                        if (!dirnameDeclared && dirnameUsed) newDeclarations.add(pathImportDeclaration)
+                        if (!requireDeclared && requireUsed) newDeclarations.add(moduleImportDeclaration)
+                        // variable
+                        if (!filenameDeclared && filenameUsed || !dirnameDeclared && dirnameUsed) newDeclarations.add(filenameDeclaration)
+                        if (!dirnameDeclared && dirnameUsed) newDeclarations.add(dirnameDeclaration)
+                        if (!requireDeclared && requireUsed) newDeclarations.add(requireDeclaration)
+                        if (newDeclarations.size > 0) {
+                            // extract shebang if exists
+                            const shebang = code.match(/^#!(.*)/)?.[0] || ''
+                            ast.body.unshift(...Array.from(newDeclarations) as any)
+                            const generated: any = escodegen.generate(ast, { sourceMapWithCode: true, directive: true })
+                            return {
+                                code: (shebang ? shebang + '\n' : '') + generated.code,
+                                map: generated.map
+                            }
+                        }
+                    }
                 }
-            }
-            return null
-        },
+                return null
+            },
+        }
     } as Plugin
 }

packages/techor/tests/shim/package.json

@@ -2,8 +2,6 @@
     "name": "@test/shim",
     "version": "1.0.0",
     "private": true,
-    "main": "dist/index.cjs",
-    "module": "dist/index.mjs",
     "files": [
         "dist"
     ]

packages/techor/tests/shim/src/dirname.ts

@@ -0,0 +1 @@
+console.log(__dirname)

packages/techor/tests/shim/src/filename.ts

@@ -0,0 +1,3 @@
+import { readFileSync } from 'fs'
+
+export const a = readFileSync(__filename, 'utf-8')

packages/techor/tests/shim/src/manual.ts

@@ -0,0 +1,9 @@
+import { dirname } from 'path'
+import { fileURLToPath } from 'url'
+import { createRequire } from 'module'
+
+const __filename = fileURLToPath(import.meta.url)
+const __dirname = dirname(__filename)
+const require = createRequire(import.meta.url)
+
+console.log(__filename, __dirname, require('./a'))

packages/techor/tests/shim/src/minify.ts

@@ -0,0 +1,4 @@
+#!/usr/bin/env node
+export const a = __filename
+export const b = __dirname
+export const c = require('./a')

packages/techor/tests/shim/src/require.ts

@@ -0,0 +1,3 @@
+export default function () {
+    return require('./filename')
+}

packages/techor/tests/shim/src/shebang.ts

@@ -0,0 +1,3 @@
+#!/usr/bin/env node
+
+console.log(__dirname)

packages/techor/tests/shim/test.ts

@@ -1,22 +1,56 @@
 import { execSync } from 'node:child_process'
 import { readFileSync } from 'node:fs'
 import { join } from 'node:path'
-import { ESM_SHIM_CODE } from '../../src/plugins/esm-shim'
 
 beforeAll(() => {
-    execSync('tsx ../../src/bin build', { cwd: __dirname })
+    execSync('tsx ../../src/bin build src/**/*.ts --formats esm', { cwd: __dirname })
+    execSync('tsx ../../src/bin build src/minify.ts --formats esm -o dist/minify.min.mjs --no-clean', { cwd: __dirname })
 })
 
-it('contains ESM shim code in ESM', () => {
-    expect(readFileSync(join(__dirname, './dist/index.mjs'), 'utf-8')).toContain(ESM_SHIM_CODE)
+it('should shim __filename in esm', () => {
+    expect(readFileSync(join(__dirname, './dist/filename.mjs'), 'utf-8'))
+        .toContain([
+            `import __url_for_shim from 'url';`,
+            `const __filename = __url_for_shim.fileURLToPath(import.meta.url);`
+        ].join('\n'))
 })
 
-it('should not contain ESM shim code in CJS', () => {
-    expect(readFileSync(join(__dirname, './dist/index.cjs'), 'utf-8')).not.toContain(ESM_SHIM_CODE)
+it('should shim __dirname in esm', () => {
+    expect(readFileSync(join(__dirname, './dist/dirname.mjs'), 'utf-8'))
+        .toContain([
+            `import __url_for_shim from 'url';`,
+            `import __path_for_shim from 'path';`,
+            `const __filename = __url_for_shim.fileURLToPath(import.meta.url);`,
+            `const __dirname = __path_for_shim.dirname(__filename);`
+        ].join('\n'))
 })
 
-it('starts with #!', () => {
-    expect(readFileSync(join(__dirname, './dist/index.cjs'), 'utf-8').startsWith('#!')).toBeTruthy()
-    expect(readFileSync(join(__dirname, './dist/index.mjs'), 'utf-8').startsWith('#!')).toBeTruthy()
+it('should shim require() in esm', () => {
+    expect(readFileSync(join(__dirname, './dist/require.mjs'), 'utf-8'))
+        .toContain([
+            `import __mod_for_shim from 'module';`,
+            `const require = __mod_for_shim.createRequire(import.meta.url);`
+        ].join('\n'))
 })
 
+it('should detect manual shim in esm', () => {
+    expect(readFileSync(join(__dirname, './dist/manual.mjs'), 'utf-8'))
+        .not.toContain([
+            `__url_for_shim`,
+            `__path_for_shim`,
+            `__mod_for_shim;`
+        ].join('\n'))
+})
+
+it('should keep the shebang in the final chunk', () => {
+    expect(readFileSync(join(__dirname, './dist/shebang.mjs'), 'utf-8')).toContain('#!/usr/bin/env node\n')
+})
+
+it('should minimize auto-filling of shims', () => {
+    expect(readFileSync(join(__dirname, './dist/minify.min.mjs'), 'utf-8')).toContain(
+        [
+            '#!/usr/bin/env node',
+            'import r from"url";import e from"path";import m from"module";let o=r.fileURLToPath(import.meta.url),t=e.dirname(o),i=m.createRequire(import.meta.url),a=o,p=t,l=i("./a");export{a,p as b,l as c};'
+        ].join('\n')
+    )
+})