fix: improve identifier detection, close #9

Author: antfu

package.json

@@ -61,6 +61,7 @@
     "eslint": "^7.32.0",
     "eslint-plugin-jest": "^24.4.0",
     "esno": "^0.9.1",
+    "fast-glob": "^3.2.7",
     "jest": "^27.0.6",
     "rimraf": "^3.0.2",
     "ts-jest": "^27.0.5",

pnpm-lock.yaml

@@ -1,7 +1,7 @@
 import { Parser as HTMLParser } from 'htmlparser2'
-import { parse, ParserOptions } from '@babel/parser'
+import { parse } from '@babel/parser'
+import { PrivateName, Expression, Statement } from '@babel/types'
 import { camelize, capitalize, isHTMLTag, isSVGTag, isVoidTag } from '@vue/shared'
-import traverse from '@babel/traverse'
 import { ParseResult, TagMeta } from './types'
 
 export function parseVueSFC(code: string, id?: string): ParseResult {
@@ -43,8 +43,12 @@ export function parseVueSFC(code: string, id?: string): ParseResult {
         Object.entries(attributes).forEach(([key, value]) => {
           if (!value)
             return
-          if (key.startsWith('v-') || key.startsWith('@') || key.startsWith(':'))
-            expressions.add(value)
+          if (key.startsWith('v-') || key.startsWith('@') || key.startsWith(':')) {
+            if (key === 'v-if')
+              expressions.add(`for (let ${value}) {}`)
+            else
+              expressions.add(`(${value})`)
+          }
           if (key === 'ref')
             identifiers.add(value)
         })
@@ -71,7 +75,7 @@ export function parseVueSFC(code: string, id?: string): ParseResult {
     ontext(text) {
       if (templateLevel > 0) {
         Array.from(text.matchAll(/\{\{(.*?)\}\}/g)).forEach(([, expression]) => {
-          expressions.add(expression)
+          expressions.add(`(${expression})`)
         })
       }
     },
@@ -102,7 +106,10 @@ export function parseVueSFC(code: string, id?: string): ParseResult {
   parser.write(code)
   parser.end()
 
-  expressions.forEach(exp => getIdentifiersFromCode(exp, identifiers))
+  expressions.forEach((exp) => {
+    const nodes = parse(exp).program.body
+    nodes.forEach(node => getIdentifiersUsage(node, identifiers))
+  })
 
   return {
     id,
@@ -115,12 +122,70 @@ export function parseVueSFC(code: string, id?: string): ParseResult {
   }
 }
 
-export function getIdentifiersFromCode(code: string, identifiers = new Set<string>(), options: ParserOptions = {}) {
-  const ast = parse(code, options) as any
-  traverse(ast, {
-    Identifier(path) {
-      identifiers.add(path.node.name)
-    },
-  })
+export function getIdentifiersDeclaration(nodes: Statement[], identifiers = new Set<string>()) {
+  for (const node of nodes) {
+    if (node.type === 'ImportDeclaration') {
+      for (const specifier of node.specifiers)
+        identifiers.add(specifier.local.name)
+    }
+    else if (node.type === 'VariableDeclaration') {
+      for (const declarator of node.declarations) {
+        // @ts-expect-error
+        identifiers.add(declarator.id.name)
+      }
+    }
+    else if (node.type === 'FunctionDeclaration') {
+      if (node.id)
+        identifiers.add(node.id.name)
+    }
+    // else {
+    //   console.log(node)
+    // }
+  }
+  return identifiers
+}
+
+export function getIdentifiersUsage(node?: Expression | PrivateName | Statement, identifiers = new Set<string>()) {
+  if (!node)
+    return identifiers
+
+  if (node.type === 'ExpressionStatement') {
+    getIdentifiersUsage(node.expression, identifiers)
+  }
+  else if (node.type === 'Identifier') {
+    identifiers.add(node.name)
+  }
+  else if (node.type === 'MemberExpression') {
+    getIdentifiersUsage(node.object, identifiers)
+  }
+  else if (node.type === 'CallExpression') {
+    // @ts-expect-error
+    getIdentifiersUsage(node.callee, identifiers)
+    node.arguments.forEach((arg) => {
+      // @ts-expect-error
+      getIdentifiersUsage(arg, identifiers)
+    })
+  }
+  else if (node.type === 'BinaryExpression' || node.type === 'LogicalExpression') {
+    getIdentifiersUsage(node.left, identifiers)
+    getIdentifiersUsage(node.right, identifiers)
+  }
+  else if (node.type === 'ForOfStatement' || node.type === 'ForInStatement') {
+    getIdentifiersUsage(node.right, identifiers)
+  }
+  else if (node.type === 'ConditionalExpression') {
+    getIdentifiersUsage(node.test, identifiers)
+    getIdentifiersUsage(node.consequent, identifiers)
+    getIdentifiersUsage(node.alternate, identifiers)
+  }
+  else if (node.type === 'ObjectExpression') {
+    node.properties.forEach((prop) => {
+      // @ts-expect-error
+      getIdentifiersUsage(prop.value, identifiers)
+    })
+  }
+  // else {
+  //   console.log(node)
+  // }
   return identifiers
 }

src/parse.ts

@@ -16,12 +16,21 @@ export function transform(sfc: string, id?: string) {
     .map(([key, value]) => value ? `${key}="${value}"` : key)
     .join(' ')
 
-  s.remove(result.script.start, result.script.end)
-  s.overwrite(
-    result.scriptSetup.start,
-    result.scriptSetup.end,
-    `<script ${attr}>\n${code}\n</script>`,
-  )
+  if (code) {
+    const block = `<script ${attr}>\n${code}\n</script>`
+
+    s.remove(result.script.start, result.script.end)
+    if (result.scriptSetup.start !== result.scriptSetup.end) {
+      s.overwrite(
+        result.scriptSetup.start,
+        result.scriptSetup.end,
+        block,
+      )
+    }
+    else {
+      s.prependLeft(0, `${block}\n`)
+    }
+  }
 
   return {
     code: s.toString(),

src/transform.ts

@@ -5,6 +5,7 @@ import traverse from '@babel/traverse'
 import generate from '@babel/generator'
 import { ParseResult } from './types'
 import { applyMacros } from './macros'
+import { getIdentifiersDeclaration } from './parse'
 
 export function transformScriptSetup(result: ParseResult) {
   if (result.script.found && result.scriptSetup.found && result.scriptSetup.attrs.lang !== result.script.attrs.lang)
@@ -26,7 +27,7 @@ export function transformScriptSetup(result: ParseResult) {
     sourceType: 'module',
     plugins,
   })
-  const scriptAst = parse(result.script.content || 'export default {}', {
+  const scriptAst = parse(result.script.content || '', {
     sourceType: 'module',
     plugins,
   })
@@ -37,12 +38,7 @@ export function transformScriptSetup(result: ParseResult) {
   const { nodes: scriptSetupBody, props } = applyMacros(nodes)
 
   // get all identifiers in `<script setup>`
-  scriptSetupAst.program.body = [...imports, ...nodes]
-  traverse(scriptSetupAst as any, {
-    Identifier(path) {
-      identifiers.add(path.node.name)
-    },
-  })
+  getIdentifiersDeclaration([...imports, ...nodes], identifiers)
 
   const returns = Array.from(identifiers).filter(i => result.template.identifiers.has(i))
   const components = Array.from(identifiers).filter(i => result.template.components.has(i)
@@ -56,10 +52,13 @@ export function transformScriptSetup(result: ParseResult) {
 
   const __sfc = t.identifier('__sfc_main')
 
+  let hasBody = false
+
   // replace `export default` with a temproray variable
   // `const __sfc_main = { ... }`
   traverse(scriptAst as any, {
     ExportDefaultDeclaration(path) {
+      hasBody = true
       const decl = path.node.declaration
       path.replaceWith(
         t.variableDeclaration('const', [
@@ -72,9 +71,22 @@ export function transformScriptSetup(result: ParseResult) {
     },
   })
 
+  // inject `const __sfc_main = {}` if `<script>` has default export
+  if (!hasBody) {
+    scriptAst.program.body.push(
+      t.variableDeclaration('const', [
+        t.variableDeclarator(
+          __sfc,
+          t.objectExpression([]),
+        ),
+      ]),
+    )
+  }
+
   // inject props function
   // `__sfc_main.props = { ... }`
   if (props) {
+    hasBody = true
     scriptAst.program.body.push(
       t.expressionStatement(
         t.assignmentExpression('=',
@@ -88,6 +100,7 @@ export function transformScriptSetup(result: ParseResult) {
   // inject setup function
   // `__sfc_main.setup = () => {}`
   if (nodes.length) {
+    hasBody = true
     const returnStatement = t.returnStatement(
       t.objectExpression(
         returns.map((i) => {
@@ -116,6 +129,7 @@ export function transformScriptSetup(result: ParseResult) {
   // inject components
   // `__sfc_main.components = Object.assign({ ... }, __sfc_main.components)`
   if (components.length) {
+    hasBody = true
     const componentsObject = t.objectExpression(
       components.map((i) => {
         const id = t.identifier(i)
@@ -139,6 +153,13 @@ export function transformScriptSetup(result: ParseResult) {
     )
   }
 
+  if (!hasBody) {
+    return {
+      ast: null,
+      code: '',
+    }
+  }
+
   // re-export
   // `export default __sfc_main`
   scriptAst.program.body.push(