From 06051c4bf207ff9ac09292b8a5a73b254608cf0e Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Mon, 23 Aug 2021 16:00:46 -0400
Subject: [PATCH] feat(experimental): support ref transform for sfc normal
 `<script>`

---
 packages/compiler-core/src/babelUtils.ts      | 12 ++-
 .../compileScriptRefSugar.spec.ts.snap        | 43 ----------
 .../compileScriptRefTransform.spec.ts.snap    | 80 +++++++++++++++++++
 ...c.ts => compileScriptRefTransform.spec.ts} | 60 +++++++++++++-
 packages/compiler-sfc/src/compileScript.ts    | 57 +++++++++++--
 .../__snapshots__/refTransform.spec.ts.snap   |  6 +-
 .../__tests__/refTransform.spec.ts            | 12 ++-
 packages/ref-transform/src/index.ts           | 30 +++----
 8 files changed, 226 insertions(+), 74 deletions(-)
 delete mode 100644 packages/compiler-sfc/__tests__/__snapshots__/compileScriptRefSugar.spec.ts.snap
 create mode 100644 packages/compiler-sfc/__tests__/__snapshots__/compileScriptRefTransform.spec.ts.snap
 rename packages/compiler-sfc/__tests__/{compileScriptRefSugar.spec.ts => compileScriptRefTransform.spec.ts} (63%)

diff --git a/packages/compiler-core/src/babelUtils.ts b/packages/compiler-core/src/babelUtils.ts
index 6faa89e261e..fe0902809b9 100644
--- a/packages/compiler-core/src/babelUtils.ts
+++ b/packages/compiler-core/src/babelUtils.ts
@@ -4,7 +4,8 @@ import {
   Node,
   Function,
   ObjectProperty,
-  BlockStatement
+  BlockStatement,
+  Program
 } from '@babel/types'
 import { walk } from 'estree-walker'
 
@@ -149,16 +150,23 @@ export function walkFunctionParams(
 }
 
 export function walkBlockDeclarations(
-  block: BlockStatement,
+  block: BlockStatement | Program,
   onIdent: (node: Identifier) => void
 ) {
   for (const stmt of block.body) {
     if (stmt.type === 'VariableDeclaration') {
+      if (stmt.declare) continue
       for (const decl of stmt.declarations) {
         for (const id of extractIdentifiers(decl.id)) {
           onIdent(id)
         }
       }
+    } else if (
+      stmt.type === 'FunctionDeclaration' ||
+      stmt.type === 'ClassDeclaration'
+    ) {
+      if (stmt.declare || !stmt.id) continue
+      onIdent(stmt.id)
     }
   }
 }
diff --git a/packages/compiler-sfc/__tests__/__snapshots__/compileScriptRefSugar.spec.ts.snap b/packages/compiler-sfc/__tests__/__snapshots__/compileScriptRefSugar.spec.ts.snap
deleted file mode 100644
index ed4605c1662..00000000000
--- a/packages/compiler-sfc/__tests__/__snapshots__/compileScriptRefSugar.spec.ts.snap
+++ /dev/null
@@ -1,43 +0,0 @@
-// Jest Snapshot v1, https://goo.gl/fbAQLP
-
-exports[`<script setup> ref sugar $ unwrapping 1`] = `
-"import { ref, shallowRef } from 'vue'
-    
-export default {
-  setup(__props, { expose }) {
-  expose()
-
-    let foo = (ref())
-    let a = (ref(1))
-    let b = (shallowRef({
-      count: 0
-    }))
-    let c = () => {}
-    let d
-    
-return { foo, a, b, c, d, ref, shallowRef }
-}
-
-}"
-`;
-
-exports[`<script setup> ref sugar $ref & $shallowRef declarations 1`] = `
-"import { ref as _ref, shallowRef as _shallowRef } from 'vue'
-
-export default {
-  setup(__props, { expose }) {
-  expose()
-
-    let foo = _ref()
-    let a = _ref(1)
-    let b = _shallowRef({
-      count: 0
-    })
-    let c = () => {}
-    let d
-    
-return { foo, a, b, c, d }
-}
-
-}"
-`;
diff --git a/packages/compiler-sfc/__tests__/__snapshots__/compileScriptRefTransform.spec.ts.snap b/packages/compiler-sfc/__tests__/__snapshots__/compileScriptRefTransform.spec.ts.snap
new file mode 100644
index 00000000000..f5674a371a9
--- /dev/null
+++ b/packages/compiler-sfc/__tests__/__snapshots__/compileScriptRefTransform.spec.ts.snap
@@ -0,0 +1,80 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`sfc ref transform $ unwrapping 1`] = `
+"import { ref, shallowRef } from 'vue'
+    
+export default {
+  setup(__props, { expose }) {
+  expose()
+
+    let foo = (ref())
+    let a = (ref(1))
+    let b = (shallowRef({
+      count: 0
+    }))
+    let c = () => {}
+    let d
+    
+return { foo, a, b, c, d, ref, shallowRef }
+}
+
+}"
+`;
+
+exports[`sfc ref transform $ref & $shallowRef declarations 1`] = `
+"import { ref as _ref, shallowRef as _shallowRef } from 'vue'
+
+export default {
+  setup(__props, { expose }) {
+  expose()
+
+    let foo = _ref()
+    let a = _ref(1)
+    let b = _shallowRef({
+      count: 0
+    })
+    let c = () => {}
+    let d
+    
+return { foo, a, b, c, d }
+}
+
+}"
+`;
+
+exports[`sfc ref transform usage in normal <script> 1`] = `
+"import { ref as _ref } from 'vue'
+
+    export default {
+      setup() {
+        let count = _ref(0)
+        const inc = () => count.value++
+        return ({ count })
+      }
+    }
+    "
+`;
+
+exports[`sfc ref transform usage with normal <script> + <script setup> 1`] = `
+"import { ref as _ref } from 'vue'
+
+    let a = _ref(0)
+    let c = _ref(0)
+    
+export default {
+  setup(__props, { expose }) {
+  expose()
+
+    let b = _ref(0)
+    let c = 0
+    function change() {
+      a.value++
+      b.value++
+      c++
+    }
+    
+return { a, c, b, change }
+}
+
+}"
+`;
diff --git a/packages/compiler-sfc/__tests__/compileScriptRefSugar.spec.ts b/packages/compiler-sfc/__tests__/compileScriptRefTransform.spec.ts
similarity index 63%
rename from packages/compiler-sfc/__tests__/compileScriptRefSugar.spec.ts
rename to packages/compiler-sfc/__tests__/compileScriptRefTransform.spec.ts
index 67c66150c72..13c011f4ab0 100644
--- a/packages/compiler-sfc/__tests__/compileScriptRefSugar.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScriptRefTransform.spec.ts
@@ -3,13 +3,13 @@ import { compileSFCScript as compile, assertCode } from './utils'
 
 // this file only tests integration with SFC - main test case for the ref
 // transform can be found in <root>/packages/ref-transform/__tests__
-describe('<script setup> ref sugar', () => {
-  function compileWithRefSugar(src: string) {
+describe('sfc ref transform', () => {
+  function compileWithRefTransform(src: string) {
     return compile(src, { refSugar: true })
   }
 
   test('$ unwrapping', () => {
-    const { content, bindings } = compileWithRefSugar(`<script setup>
+    const { content, bindings } = compileWithRefTransform(`<script setup>
     import { ref, shallowRef } from 'vue'
     let foo = $(ref())
     let a = $(ref(1))
@@ -46,7 +46,7 @@ describe('<script setup> ref sugar', () => {
   })
 
   test('$ref & $shallowRef declarations', () => {
-    const { content, bindings } = compileWithRefSugar(`<script setup>
+    const { content, bindings } = compileWithRefTransform(`<script setup>
     let foo = $ref()
     let a = $ref(1)
     let b = $shallowRef({
@@ -81,6 +81,58 @@ describe('<script setup> ref sugar', () => {
     })
   })
 
+  test('usage in normal <script>', () => {
+    const { content } = compileWithRefTransform(`<script>
+    export default {
+      setup() {
+        let count = $ref(0)
+        const inc = () => count++
+        return $$({ count })
+      }
+    }
+    </script>`)
+    expect(content).not.toMatch(`$ref(0)`)
+    expect(content).toMatch(`import { ref as _ref } from 'vue'`)
+    expect(content).toMatch(`let count = _ref(0)`)
+    expect(content).toMatch(`count.value++`)
+    expect(content).toMatch(`return ({ count })`)
+    assertCode(content)
+  })
+
+  test('usage with normal <script> + <script setup>', () => {
+    const { content, bindings } = compileWithRefTransform(`<script>
+    let a = $ref(0)
+    let c = $ref(0)
+    </script>
+    <script setup>
+    let b = $ref(0)
+    let c = 0
+    function change() {
+      a++
+      b++
+      c++
+    }
+    </script>`)
+    // should dedupe helper imports
+    expect(content).toMatch(`import { ref as _ref } from 'vue'`)
+
+    expect(content).toMatch(`let a = _ref(0)`)
+    expect(content).toMatch(`let b = _ref(0)`)
+
+    // root level ref binding declared in <script> should be inherited in <script setup>
+    expect(content).toMatch(`a.value++`)
+    expect(content).toMatch(`b.value++`)
+    // c shadowed
+    expect(content).toMatch(`c++`)
+    assertCode(content)
+    expect(bindings).toStrictEqual({
+      a: BindingTypes.SETUP_REF,
+      b: BindingTypes.SETUP_REF,
+      c: BindingTypes.SETUP_REF,
+      change: BindingTypes.SETUP_CONST
+    })
+  })
+
   describe('errors', () => {
     test('defineProps/Emit() referencing ref declarations', () => {
       expect(() =>
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index d71aa90b6a8..55e5e420511 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -165,12 +165,34 @@ export function compileScript(
       return script
     }
     try {
-      const scriptAst = _parse(script.content, {
+      let content = script.content
+      let map = script.map
+      const scriptAst = _parse(content, {
         plugins,
         sourceType: 'module'
-      }).program.body
-      const bindings = analyzeScriptBindings(scriptAst)
-      let content = script.content
+      }).program
+      const bindings = analyzeScriptBindings(scriptAst.body)
+      if (enableRefTransform && shouldTransformRef(content)) {
+        const s = new MagicString(source)
+        const startOffset = script.loc.start.offset
+        const endOffset = script.loc.end.offset
+        const { importedHelpers } = transformRefAST(scriptAst, s, startOffset)
+        if (importedHelpers.length) {
+          s.prepend(
+            `import { ${importedHelpers
+              .map(h => `${h} as _${h}`)
+              .join(', ')} } from 'vue'\n`
+          )
+        }
+        s.remove(0, startOffset)
+        s.remove(endOffset, source.length)
+        content = s.toString()
+        map = s.generateMap({
+          source: filename,
+          hires: true,
+          includeContent: true
+        }) as unknown as RawSourceMap
+      }
       if (cssVars.length) {
         content = rewriteDefault(content, `__default__`, plugins)
         content += genNormalScriptCssVarsCode(
@@ -184,8 +206,9 @@ export function compileScript(
       return {
         ...script,
         content,
+        map,
         bindings,
-        scriptAst
+        scriptAst: scriptAst.body
       }
     } catch (e) {
       // silently fallback if parse fails since user may be using custom
@@ -629,6 +652,23 @@ export function compileScript(
         walkDeclaration(node, setupBindings, userImportAlias)
       }
     }
+
+    // apply ref transform
+    if (enableRefTransform && shouldTransformRef(script.content)) {
+      warnExperimental(
+        `ref sugar`,
+        `https://github.com/vuejs/rfcs/discussions/369`
+      )
+      const { rootVars, importedHelpers } = transformRefAST(
+        scriptAst,
+        s,
+        scriptStartOffset!
+      )
+      refBindings = rootVars
+      for (const h of importedHelpers) {
+        helperImports.add(h)
+      }
+    }
   }
 
   // 2. parse <script setup> and  walk over top level statements
@@ -862,7 +902,7 @@ export function compileScript(
   }
 
   // 3. Apply ref sugar transform
-  if (enableRefTransform && shouldTransformRef(source)) {
+  if (enableRefTransform && shouldTransformRef(scriptSetup.content)) {
     warnExperimental(
       `ref sugar`,
       `https://github.com/vuejs/rfcs/discussions/369`
@@ -870,9 +910,10 @@ export function compileScript(
     const { rootVars, importedHelpers } = transformRefAST(
       scriptSetupAst,
       s,
-      startOffset
+      startOffset,
+      refBindings
     )
-    refBindings = rootVars
+    refBindings = refBindings ? [...refBindings, ...rootVars] : rootVars
     for (const h of importedHelpers) {
       helperImports.add(h)
     }
diff --git a/packages/ref-transform/__tests__/__snapshots__/refTransform.spec.ts.snap b/packages/ref-transform/__tests__/__snapshots__/refTransform.spec.ts.snap
index 54044f47c3c..5ad3f93bb18 100644
--- a/packages/ref-transform/__tests__/__snapshots__/refTransform.spec.ts.snap
+++ b/packages/ref-transform/__tests__/__snapshots__/refTransform.spec.ts.snap
@@ -136,6 +136,9 @@ exports[`nested scopes 1`] = `
     b.value++ // outer b
     c++ // outer c
 
+    let bar = _ref(0)
+    bar.value++ // outer bar
+
     function foo({ a }) {
       a++ // inner a
       b.value++ // inner b
@@ -143,10 +146,11 @@ exports[`nested scopes 1`] = `
       c.value++ // inner c
       let d = _ref(0)
 
-      const bar = (c) => {
+      function bar(c) {
         c++ // nested c
         d.value++ // nested d
       }
+      bar() // inner bar
 
       if (true) {
         let a = _ref(0)
diff --git a/packages/ref-transform/__tests__/refTransform.spec.ts b/packages/ref-transform/__tests__/refTransform.spec.ts
index 8164f977910..462fee5fd90 100644
--- a/packages/ref-transform/__tests__/refTransform.spec.ts
+++ b/packages/ref-transform/__tests__/refTransform.spec.ts
@@ -279,6 +279,9 @@ test('nested scopes', () => {
     b++ // outer b
     c++ // outer c
 
+    let bar = $ref(0)
+    bar++ // outer bar
+
     function foo({ a }) {
       a++ // inner a
       b++ // inner b
@@ -286,10 +289,11 @@ test('nested scopes', () => {
       c++ // inner c
       let d = $ref(0)
 
-      const bar = (c) => {
+      function bar(c) {
         c++ // nested c
         d++ // nested d
       }
+      bar() // inner bar
 
       if (true) {
         let a = $ref(0)
@@ -299,7 +303,7 @@ test('nested scopes', () => {
       return $$({ a, b, c, d })
     }
     `)
-  expect(rootVars).toStrictEqual(['a', 'b'])
+  expect(rootVars).toStrictEqual(['a', 'b', 'bar'])
 
   expect(code).toMatch('a.value++ // outer a')
   expect(code).toMatch('b.value++ // outer b')
@@ -314,6 +318,10 @@ test('nested scopes', () => {
 
   expect(code).toMatch(`a.value++ // if block a`) // if block
 
+  expect(code).toMatch(`bar.value++ // outer bar`)
+  // inner bar shadowed by function declaration
+  expect(code).toMatch(`bar() // inner bar`)
+
   expect(code).toMatch(`return ({ a, b, c, d })`)
   assertCode(code)
 })
diff --git a/packages/ref-transform/src/index.ts b/packages/ref-transform/src/index.ts
index 73c5b086a77..87f1e5356db 100644
--- a/packages/ref-transform/src/index.ts
+++ b/packages/ref-transform/src/index.ts
@@ -1,7 +1,6 @@
 import {
   Node,
   Identifier,
-  VariableDeclarator,
   BlockStatement,
   CallExpression,
   ObjectPattern,
@@ -30,14 +29,6 @@ export function shouldTransform(src: string): boolean {
   return transformCheckRE.test(src)
 }
 
-export interface ReactiveDeclarator {
-  node: VariableDeclarator
-  statement: VariableDeclaration
-  ids: Identifier[]
-  isPattern: boolean
-  isRoot: boolean
-}
-
 type Scope = Record<string, boolean>
 
 export interface RefTransformOptions {
@@ -105,18 +96,26 @@ export function transform(
 export function transformAST(
   ast: Node,
   s: MagicString,
-  offset = 0
+  offset = 0,
+  knownRootVars?: string[]
 ): {
   rootVars: string[]
   importedHelpers: string[]
 } {
   const importedHelpers = new Set<string>()
   const blockStack: BlockStatement[] = []
+  let currentBlock: BlockStatement | null = null
   const rootScope: Scope = {}
   const blockToScopeMap = new WeakMap<BlockStatement, Scope>()
   const excludedIds = new Set<Identifier>()
   const parentStack: Node[] = []
 
+  if (knownRootVars) {
+    for (const key of knownRootVars) {
+      rootScope[key] = true
+    }
+  }
+
   const error = (msg: string, node: Node) => {
     const e = new Error(msg)
     ;(e as any).node = node
@@ -130,7 +129,6 @@ export function transformAST(
 
   const registerBinding = (id: Identifier, isRef = false) => {
     excludedIds.add(id)
-    const currentBlock = blockStack[blockStack.length - 1]
     if (currentBlock) {
       const currentScope = blockToScopeMap.get(currentBlock)
       if (!currentScope) {
@@ -145,13 +143,16 @@ export function transformAST(
 
   const registerRefBinding = (id: Identifier) => registerBinding(id, true)
 
+  if (ast.type === 'Program') {
+    walkBlockDeclarations(ast, registerBinding)
+  }
+
   // 1st pass: detect macro callsites and register ref bindings
   ;(walk as any)(ast, {
     enter(node: Node, parent?: Node) {
       parent && parentStack.push(parent)
-
       if (node.type === 'BlockStatement') {
-        blockStack.push(node)
+        blockStack.push((currentBlock = node))
         walkBlockDeclarations(node, registerBinding)
         if (parent && isFunctionType(parent)) {
           walkFunctionParams(parent, registerBinding)
@@ -213,6 +214,7 @@ export function transformAST(
       parent && parentStack.pop()
       if (node.type === 'BlockStatement') {
         blockStack.pop()
+        currentBlock = blockStack[blockStack.length - 1] || null
       }
     }
   })
@@ -356,7 +358,7 @@ export function transformAST(
   }
 
   return {
-    rootVars: Object.keys(rootScope),
+    rootVars: Object.keys(rootScope).filter(key => rootScope[key]),
     importedHelpers: [...importedHelpers]
   }
 }