From 6861d2380b87566cbc479e065fb67dfeb8c86c46 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Tue, 8 Nov 2022 17:47:47 +0800
Subject: [PATCH] refactor(compiler-sfc): optimize import alias check for
 binding analysis

---
 .../__snapshots__/compileScript.spec.ts.snap  |  33 +--
 .../__tests__/compileScript.spec.ts           |   4 +-
 packages/compiler-sfc/src/compileScript.ts    | 214 +++++++++---------
 3 files changed, 127 insertions(+), 124 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
index 209bbe1d518..ec94872c560 100644
--- a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
+++ b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
@@ -44,9 +44,9 @@ return { a }
 `;
 
 exports[`SFC compile <script setup> <script> after <script setup> the script content not end with \`\\n\` 1`] = `
-"const n = 1
-import { x } from './x'
-    
+"import { x } from './x'
+    const n = 1
+
 export default {
   setup(__props, { expose }) {
   expose();
@@ -78,11 +78,12 @@ return { n, x }
 `;
 
 exports[`SFC compile <script setup> <script> and <script setup> co-usage script setup first 1`] = `
-"export const n = 1
-      const __default__ = {}
+"import { x } from './x'
       
-import { x } from './x'
+      export const n = 1
+      const __default__ = {}
       
+
 export default /*#__PURE__*/Object.assign(__default__, {
   setup(__props, { expose }) {
   expose();
@@ -97,13 +98,13 @@ return { n, x }
 
 exports[`SFC compile <script setup> <script> and <script setup> co-usage script setup first, lang="ts", script block content export default 1`] = `
 "import { defineComponent as _defineComponent } from 'vue'
-
+import { x } from './x'
+      
       const __default__ = {
         name: \\"test\\"
       }
       
-import { x } from './x'
-      
+
 export default /*#__PURE__*/_defineComponent({
   ...__default__,
   setup(__props, { expose }) {
@@ -118,14 +119,15 @@ return { x }
 `;
 
 exports[`SFC compile <script setup> <script> and <script setup> co-usage script setup first, named default export 1`] = `
-"export const n = 1
+"import { x } from './x'
+      
+      export const n = 1
       const def = {}
       
       
 const __default__ = def
 
-import { x } from './x'
-      
+
 export default /*#__PURE__*/Object.assign(__default__, {
   setup(__props, { expose }) {
   expose();
@@ -1260,14 +1262,15 @@ return () => {}
 `;
 
 exports[`SFC compile <script setup> should expose top level declarations 1`] = `
-"import { xx } from './x'
+"import { x } from './x'
+      
+      import { xx } from './x'
       let aa = 1
       const bb = 2
       function cc() {}
       class dd {}
       
-import { x } from './x'
-      
+
 export default {
   setup(__props, { expose }) {
   expose();
diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index 96160e272d5..4e3e36f7431 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -387,8 +387,8 @@ defineExpose({ foo: 123 })
         expect(bindings).toStrictEqual({
           _reactive: BindingTypes.SETUP_MAYBE_REF,
           _ref: BindingTypes.SETUP_MAYBE_REF,
-          foo: BindingTypes.SETUP_REF,
-          bar: BindingTypes.SETUP_REACTIVE_CONST
+          foo: BindingTypes.SETUP_MAYBE_REF,
+          bar: BindingTypes.SETUP_MAYBE_REF
         })
       })
 
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index ccf670725ec..aba5a22ec8d 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -898,10 +898,10 @@ export function compileScript(
     }
   }
 
-  // 1. process normal <script> first if it exists
-  let scriptAst: Program | undefined
-  if (script) {
-    scriptAst = parse(
+  // 0. parse both <script> and <script setup> blocks
+  const scriptAst =
+    script &&
+    parse(
       script.content,
       {
         plugins,
@@ -910,7 +910,21 @@ export function compileScript(
       scriptStartOffset!
     )
 
-    // walk import declarations first
+  const scriptSetupAst = parse(
+    scriptSetup.content,
+    {
+      plugins: [
+        ...plugins,
+        // allow top level await but only inside <script setup>
+        'topLevelAwait'
+      ],
+      sourceType: 'module'
+    },
+    startOffset
+  )
+
+  // 1.1 walk import delcarations of <script>
+  if (scriptAst) {
     for (const node of scriptAst.body) {
       if (node.type === 'ImportDeclaration') {
         // record imports for dedupe
@@ -932,7 +946,88 @@ export function compileScript(
         }
       }
     }
+  }
+
+  // 1.2 walk import declarations of <script setup>
+  for (const node of scriptSetupAst.body) {
+    if (node.type === 'ImportDeclaration') {
+      // import declarations are moved to top
+      hoistNode(node)
+
+      // dedupe imports
+      let removed = 0
+      const removeSpecifier = (i: number) => {
+        const removeLeft = i > removed
+        removed++
+        const current = node.specifiers[i]
+        const next = node.specifiers[i + 1]
+        s.remove(
+          removeLeft
+            ? node.specifiers[i - 1].end! + startOffset
+            : current.start! + startOffset,
+          next && !removeLeft
+            ? next.start! + startOffset
+            : current.end! + startOffset
+        )
+      }
+
+      for (let i = 0; i < node.specifiers.length; i++) {
+        const specifier = node.specifiers[i]
+        const local = specifier.local.name
+        let imported =
+          specifier.type === 'ImportSpecifier' &&
+          specifier.imported.type === 'Identifier' &&
+          specifier.imported.name
+        if (specifier.type === 'ImportNamespaceSpecifier') {
+          imported = '*'
+        }
+        const source = node.source.value
+        const existing = userImports[local]
+        if (
+          source === 'vue' &&
+          (imported === DEFINE_PROPS ||
+            imported === DEFINE_EMITS ||
+            imported === DEFINE_EXPOSE)
+        ) {
+          warnOnce(
+            `\`${imported}\` is a compiler macro and no longer needs to be imported.`
+          )
+          removeSpecifier(i)
+        } else if (existing) {
+          if (existing.source === source && existing.imported === imported) {
+            // already imported in <script setup>, dedupe
+            removeSpecifier(i)
+          } else {
+            error(`different imports aliased to same local name.`, specifier)
+          }
+        } else {
+          registerUserImport(
+            source,
+            local,
+            imported,
+            node.importKind === 'type' ||
+              (specifier.type === 'ImportSpecifier' &&
+                specifier.importKind === 'type'),
+            true,
+            !options.inlineTemplate
+          )
+        }
+      }
+      if (node.specifiers.length && removed === node.specifiers.length) {
+        s.remove(node.start! + startOffset, node.end! + startOffset)
+      }
+    }
+  }
+
+  // 1.3 resolve possible user import alias of `ref` and `reactive`
+  const vueImportAliases: Record<string, string> = {}
+  for (const key in userImports) {
+    const { source, imported, local } = userImports[key]
+    if (source === 'vue') vueImportAliases[imported] = local
+  }
 
+  // 2.1 process normal <script> body
+  if (script && scriptAst) {
     for (const node of scriptAst.body) {
       if (node.type === 'ExportDefaultDeclaration') {
         // export default
@@ -1011,7 +1106,7 @@ export function compileScript(
           }
         }
         if (node.declaration) {
-          walkDeclaration(node.declaration, scriptBindings, userImports)
+          walkDeclaration(node.declaration, scriptBindings, vueImportAliases)
         }
       } else if (
         (node.type === 'VariableDeclaration' ||
@@ -1020,7 +1115,7 @@ export function compileScript(
           node.type === 'TSEnumDeclaration') &&
         !node.declare
       ) {
-        walkDeclaration(node, scriptBindings, userImports)
+        walkDeclaration(node, scriptBindings, vueImportAliases)
       }
     }
 
@@ -1049,94 +1144,8 @@ export function compileScript(
     }
   }
 
-  // 2. parse <script setup> and  walk over top level statements
-  const scriptSetupAst = parse(
-    scriptSetup.content,
-    {
-      plugins: [
-        ...plugins,
-        // allow top level await but only inside <script setup>
-        'topLevelAwait'
-      ],
-      sourceType: 'module'
-    },
-    startOffset
-  )
-
+  // 2.2 process <script setup> body
   for (const node of scriptSetupAst.body) {
-    if (node.type === 'ImportDeclaration') {
-      // import declarations are moved to top
-      hoistNode(node)
-
-      // dedupe imports
-      let removed = 0
-      const removeSpecifier = (i: number) => {
-        const removeLeft = i > removed
-        removed++
-        const current = node.specifiers[i]
-        const next = node.specifiers[i + 1]
-        s.remove(
-          removeLeft
-            ? node.specifiers[i - 1].end! + startOffset
-            : current.start! + startOffset,
-          next && !removeLeft
-            ? next.start! + startOffset
-            : current.end! + startOffset
-        )
-      }
-
-      for (let i = 0; i < node.specifiers.length; i++) {
-        const specifier = node.specifiers[i]
-        const local = specifier.local.name
-        let imported =
-          specifier.type === 'ImportSpecifier' &&
-          specifier.imported.type === 'Identifier' &&
-          specifier.imported.name
-        if (specifier.type === 'ImportNamespaceSpecifier') {
-          imported = '*'
-        }
-        const source = node.source.value
-        const existing = userImports[local]
-        if (
-          source === 'vue' &&
-          (imported === DEFINE_PROPS ||
-            imported === DEFINE_EMITS ||
-            imported === DEFINE_EXPOSE)
-        ) {
-          warnOnce(
-            `\`${imported}\` is a compiler macro and no longer needs to be imported.`
-          )
-          removeSpecifier(i)
-        } else if (existing) {
-          if (existing.source === source && existing.imported === imported) {
-            // already imported in <script setup>, dedupe
-            removeSpecifier(i)
-          } else {
-            error(`different imports aliased to same local name.`, specifier)
-          }
-        } else {
-          registerUserImport(
-            source,
-            local,
-            imported,
-            node.importKind === 'type' ||
-              (specifier.type === 'ImportSpecifier' &&
-                specifier.importKind === 'type'),
-            true,
-            !options.inlineTemplate
-          )
-        }
-      }
-      if (node.specifiers.length && removed === node.specifiers.length) {
-        s.remove(node.start! + startOffset, node.end! + startOffset)
-      }
-    }
-  }
-
-  for (const node of scriptSetupAst.body) {
-    // already processed
-    if (node.type === 'ImportDeclaration') continue
-
     // (Dropped) `ref: x` bindings
     // TODO remove when out of experimental
     if (
@@ -1210,7 +1219,7 @@ export function compileScript(
         node.type === 'ClassDeclaration') &&
       !node.declare
     ) {
-      walkDeclaration(node, setupBindings, userImports)
+      walkDeclaration(node, setupBindings, vueImportAliases)
     }
 
     // walk statements & named exports / variable declarations for top level
@@ -1665,17 +1674,8 @@ function registerBinding(
 function walkDeclaration(
   node: Declaration,
   bindings: Record<string, BindingTypes>,
-  userImports: Record<string, ImportBinding>
+  userImportAliases: Record<string, string>
 ) {
-  function getUserBinding(name: string) {
-    const binding = Object.values(userImports).find(
-      binding => binding.source === 'vue' && binding.imported === name
-    )
-    if (binding) return binding.local
-    else if (!userImports[name]) return name
-    return undefined
-  }
-
   if (node.type === 'VariableDeclaration') {
     const isConst = node.kind === 'const'
     // export const foo = ...
@@ -1689,7 +1689,7 @@ function walkDeclaration(
       )
       if (id.type === 'Identifier') {
         let bindingType
-        const userReactiveBinding = getUserBinding('reactive')
+        const userReactiveBinding = userImportAliases['reactive']
         if (isCallOf(init, userReactiveBinding)) {
           // treat reactive() calls as let since it's meant to be mutable
           bindingType = isConst
@@ -1705,7 +1705,7 @@ function walkDeclaration(
             ? BindingTypes.SETUP_REACTIVE_CONST
             : BindingTypes.SETUP_CONST
         } else if (isConst) {
-          if (isCallOf(init, getUserBinding('ref'))) {
+          if (isCallOf(init, userImportAliases['ref'])) {
             bindingType = BindingTypes.SETUP_REF
           } else {
             bindingType = BindingTypes.SETUP_MAYBE_REF