workflow(sfc-playground): improve module rewrite

packages/sfc-playground/src/output/moduleCompiler.ts

@@ -1,12 +1,23 @@
 import { store, MAIN_FILE, SANDBOX_VUE_URL, File } from '../store'
-import { babelParse, MagicString, walk } from '@vue/compiler-sfc'
+import {
+  babelParse,
+  MagicString,
+  walk,
+  walkIdentifiers
+} from '@vue/compiler-sfc'
 import { babelParserDefaultPlugins } from '@vue/shared'
-import { Identifier, Node } from '@babel/types'
+import { ExportSpecifier, Identifier, Node, ObjectProperty } from '@babel/types'
 
 export function compileModulesForPreview() {
   return processFile(store.files[MAIN_FILE]).reverse()
 }
 
+const modulesKey = `__modules__`
+const exportKey = `__export__`
+const dynamicImportKey = `__dynamic_import__`
+const moduleKey = `__module__`
+
+// similar logic with Vite's SSR transform, except this is targeting the browser
 function processFile(file: File, seen = new Set<File>()) {
   if (seen.has(file)) {
     return []
@@ -14,133 +25,191 @@ function processFile(file: File, seen = new Set<File>()) {
   seen.add(file)
 
   const { js, css } = file.compiled
+
+  const s = new MagicString(js)
+
   const ast = babelParse(js, {
     sourceFilename: file.filename,
     sourceType: 'module',
     plugins: [...babelParserDefaultPlugins]
   }).program.body
 
+  const idToImportMap = new Map<string, string>()
+  const declaredConst = new Set<string>()
   const importedFiles = new Set<string>()
   const importToIdMap = new Map<string, string>()
 
-  const s = new MagicString(js)
+  function defineImport(node: Node, source: string) {
-
-  function registerImport(source: string) {
     const filename = source.replace(/^\.\/+/, '')
     if (!(filename in store.files)) {
       throw new Error(`File "${filename}" does not exist.`)
     }
     if (importedFiles.has(filename)) {
-      return importToIdMap.get(filename)
+      return importToIdMap.get(filename)!
     }
     importedFiles.add(filename)
     const id = `__import_${importedFiles.size}__`
     importToIdMap.set(filename, id)
-    s.prepend(`const ${id} = __modules__[${JSON.stringify(filename)}]\n`)
+    s.appendLeft(
+      node.start!,
+      `const ${id} = ${modulesKey}[${JSON.stringify(filename)}]\n`
+    )
     return id
   }
 
+  function defineExport(name: string, local = name) {
+    s.append(`\n${exportKey}(${moduleKey}, "${name}", () => ${local})`)
+  }
+
+  // 0. instantiate module
   s.prepend(
-    `const mod = __modules__[${JSON.stringify(
+    `const ${moduleKey} = __modules__[${JSON.stringify(
       file.filename
-    )}] = Object.create(null)\n\n`
+    )}] = { [Symbol.toStringTag]: "Module" }\n\n`
   )
 
+  // 1. check all import statements and record id -> importName map
   for (const node of ast) {
+    // import foo from 'foo' --> foo -> __import_foo__.default
+    // import { baz } from 'foo' --> baz -> __import_foo__.baz
+    // import * as ok from 'foo' --> ok -> __import_foo__
     if (node.type === 'ImportDeclaration') {
       const source = node.source.value
-      if (source === 'vue') {
+      if (source.startsWith('./')) {
-        // rewrite Vue imports
+        const importId = defineImport(node, node.source.value)
-        s.overwrite(
-          node.source.start!,
-          node.source.end!,
-          `"${SANDBOX_VUE_URL}"`
-        )
-      } else if (source.startsWith('./')) {
-        // rewrite the import to retrieve the import from global registry
-        s.remove(node.start!, node.end!)
-
-        const id = registerImport(source)
-
         for (const spec of node.specifiers) {
-          if (spec.type === 'ImportDefaultSpecifier') {
+          if (spec.type === 'ImportSpecifier') {
-            s.prependRight(
+            idToImportMap.set(
-              node.start!,
+              spec.local.name,
-              `const ${spec.local.name} = ${id}.default\n`
+              `${importId}.${(spec.imported as Identifier).name}`
-            )
-          } else if (spec.type === 'ImportSpecifier') {
-            s.prependRight(
-              node.start!,
-              `const ${spec.local.name} = ${id}.${
-                (spec.imported as Identifier).name
-              }\n`
             )
+          } else if (spec.type === 'ImportDefaultSpecifier') {
+            idToImportMap.set(spec.local.name, `${importId}.default`)
           } else {
-            // namespace import
+            // namespace specifier
-            s.prependRight(node.start!, `const ${spec.local.name} = ${id}`)
+            idToImportMap.set(spec.local.name, importId)
           }
         }
+        s.remove(node.start!, node.end!)
+      } else {
+        if (source === 'vue') {
+          // rewrite Vue imports
+          s.overwrite(
+            node.source.start!,
+            node.source.end!,
+            `"${SANDBOX_VUE_URL}"`
+          )
+        }
       }
     }
+  }
 
-    if (node.type === 'ExportDefaultDeclaration') {
+  // 2. check all export statements and define exports
-      // export default -> mod.default = ...
+  for (const node of ast) {
-      s.overwrite(node.start!, node.declaration.start!, 'mod.default = ')
+    // named exports
-    }
-
     if (node.type === 'ExportNamedDeclaration') {
-      if (node.source) {
+      if (node.declaration) {
-        // export { foo } from '...' -> mode.foo = __import_x__.foo
-        const id = registerImport(node.source.value)
-        let code = ``
-        for (const spec of node.specifiers) {
-          if (spec.type === 'ExportSpecifier') {
-            code += `mod.${(spec.exported as Identifier).name} = ${id}.${
-              spec.local.name
-            }\n`
-          }
-        }
-        s.overwrite(node.start!, node.end!, code)
-      } else if (node.declaration) {
         if (
           node.declaration.type === 'FunctionDeclaration' ||
           node.declaration.type === 'ClassDeclaration'
         ) {
           // export function foo() {}
-          const name = node.declaration.id!.name
+          defineExport(node.declaration.id!.name)
-          s.appendLeft(node.end!, `\nmod.${name} = ${name}\n`)
         } else if (node.declaration.type === 'VariableDeclaration') {
           // export const foo = 1, bar = 2
           for (const decl of node.declaration.declarations) {
-            for (const { name } of extractIdentifiers(decl.id)) {
+            const names = extractNames(decl.id as any)
-              s.appendLeft(node.end!, `\nmod.${name} = ${name}`)
+            for (const name of names) {
+              defineExport(name)
             }
           }
         }
         s.remove(node.start!, node.declaration.start!)
-      } else {
+      } else if (node.source) {
-        let code = ``
+        // export { foo, bar } from './foo'
+        const importId = defineImport(node, node.source.value)
         for (const spec of node.specifiers) {
-          if (spec.type === 'ExportSpecifier') {
+          defineExport(
-            code += `mod.${(spec.exported as Identifier).name} = ${
+            (spec.exported as Identifier).name,
-              spec.local.name
+            `${importId}.${(spec as ExportSpecifier).local.name}`
-            }\n`
+          )
-          }
         }
-        s.overwrite(node.start!, node.end!, code)
+        s.remove(node.start!, node.end!)
+      } else {
+        // export { foo, bar }
+        for (const spec of node.specifiers) {
+          const local = (spec as ExportSpecifier).local.name
+          const binding = idToImportMap.get(local)
+          defineExport((spec.exported as Identifier).name, binding || local)
+        }
+        s.remove(node.start!, node.end!)
       }
     }
 
+    // default export
+    if (node.type === 'ExportDefaultDeclaration') {
+      s.overwrite(node.start!, node.start! + 14, `${moduleKey}.default =`)
+    }
+
+    // export * from './foo'
     if (node.type === 'ExportAllDeclaration') {
-      const id = registerImport(node.source.value)
+      const importId = defineImport(node, node.source.value)
-      s.overwrite(node.start!, node.end!, `Object.assign(mod, ${id})`)
+      s.remove(node.start!, node.end!)
+      s.append(`\nfor (const key in ${importId}) {
+        if (key !== 'default') {
+          ${exportKey}(${moduleKey}, key, () => ${importId}[key])
+        }
+      }`)
     }
   }
 
-  // dynamic import
+  // 3. convert references to import bindings
-  walk(ast as any, {
+  for (const node of ast) {
-    enter(node) {
+    if (node.type === 'ImportDeclaration') continue
-      if (node.type === 'ImportExpression') {
+    walkIdentifiers(node, (id, parent, parentStack) => {
+      const binding = idToImportMap.get(id.name)
+      if (!binding) {
+        return
+      }
+      if (isStaticProperty(parent) && parent.shorthand) {
+        // let binding used in a property shorthand
+        // { foo } -> { foo: __import_x__.foo }
+        // skip for destructure patterns
+        if (
+          !(parent as any).inPattern ||
+          isInDestructureAssignment(parent, parentStack)
+        ) {
+          s.appendLeft(id.end!, `: ${binding}`)
+        }
+      } else if (
+        parent.type === 'ClassDeclaration' &&
+        id === parent.superClass
+      ) {
+        if (!declaredConst.has(id.name)) {
+          declaredConst.add(id.name)
+          // locate the top-most node containing the class declaration
+          const topNode = parentStack[1]
+          s.prependRight(topNode.start!, `const ${id.name} = ${binding};\n`)
+        }
+      } else {
+        s.overwrite(id.start!, id.end!, binding)
+      }
+    })
+  }
+
+  // 4. convert dynamic imports
+  ;(walk as any)(ast, {
+    enter(node: Node, parent: Node) {
+      if (node.type === 'Import' && parent.type === 'CallExpression') {
+        const arg = parent.arguments[0]
+        if (arg.type === 'StringLiteral' && arg.value.startsWith('./')) {
+          s.overwrite(node.start!, node.start! + 6, dynamicImportKey)
+          s.overwrite(
+            arg.start!,
+            arg.end!,
+            JSON.stringify(arg.value.replace(/^\.\/+/, ''))
+          )
+        }
       }
     }
   })
@@ -161,6 +230,13 @@ function processFile(file: File, seen = new Set<File>()) {
   return processed
 }
 
+const isStaticProperty = (node: Node): node is ObjectProperty =>
+  node.type === 'ObjectProperty' && !node.computed
+
+function extractNames(param: Node): string[] {
+  return extractIdentifiers(param).map(id => id.name)
+}
+
 function extractIdentifiers(
   param: Node,
   nodes: Identifier[] = []
@@ -205,3 +281,21 @@ function extractIdentifiers(
 
   return nodes
 }
+
+function isInDestructureAssignment(parent: Node, parentStack: Node[]): boolean {
+  if (
+    parent &&
+    (parent.type === 'ObjectProperty' || parent.type === 'ArrayPattern')
+  ) {
+    let i = parentStack.length
+    while (i--) {
+      const p = parentStack[i]
+      if (p.type === 'AssignmentExpression') {
+        return true
+      } else if (p.type !== 'ObjectProperty' && !p.type.endsWith('Pattern')) {
+        break
+      }
+    }
+  }
+  return false
+}

packages/sfc-playground/src/output/srcdoc.html

@@ -11,6 +11,20 @@
 		<script type="module">
 			let scriptEl
 
+			window.__modules__ = {}
+
+			window.__export__ = (mod, key, get) => {
+				Object.defineProperty(mod, key, {
+					enumerable: true,
+					configurable: true,
+					get
+				})
+			}
+
+			window.__dynamic_import__ = key => {
+				return Promise.resolve(window.__modules__[key])
+			}
+
 			function handle_message(ev) {
 				let { action, cmd_id } = ev.data;
 				const send_message = (payload) => parent.postMessage( { ...payload }, ev.origin);