feat: use @babel/traverse to get identifiers

Author: xiaoxiangmoe

src/core/babel.ts

@@ -1,7 +1,9 @@
 import * as babel from '@babel/core'
 import { parse, parseExpression } from '@babel/parser'
 import g from '@babel/generator'
+import * as babel_traverse from '@babel/traverse'
 
 export const t: typeof babel['types'] = ((babel as any).default || babel).types
 export const generate: typeof g = ((g as any).default || g)
+export const traverse = ((babel_traverse as any)?.default?.default as null) ?? babel_traverse?.default ?? babel_traverse
 export { parseExpression, parse }

src/core/identifiers.ts

@@ -1,59 +1,41 @@
-import type { Expression, Node, PrivateName, SpreadElement, Statement, TSType } from '@babel/types'
+import type {
+  Expression,
+  File,
+  PrivateName,
+  SpreadElement,
+  Statement,
+  TSType,
+} from '@babel/types'
+import type { ParseResult } from '@babel/parser'
+import { t, traverse } from './babel'
 
-export function getIdentifierDeclarations(nodes: Statement[], identifiers = new Set<string>()) {
-  for (let node of nodes) {
-    if (node.type === 'ExportNamedDeclaration') {
-      node = node.declaration!
-      if (!node)
-        continue
-    }
-    if (node.type === 'ImportDeclaration') {
-      for (const specifier of node.specifiers)
-        identifiers.add(specifier.local.name)
-    }
-    else if (node.type === 'VariableDeclaration') {
-      function handleVariableId(node: Node) {
-        if (node.type === 'Identifier') {
-          identifiers.add(node.name)
-        }
-        else if (node.type === 'ObjectPattern') {
-          for (const property of node.properties) {
-            if (property.type === 'ObjectProperty')
-              handleVariableId(property.value)
-            else if (property.type === 'RestElement' && property.argument.type === 'Identifier')
-              identifiers.add(property.argument.name)
-          }
-        }
-        else if (node.type === 'ArrayPattern') {
-          for (const element of node.elements) {
-            if (element?.type === 'Identifier')
-              identifiers.add(element.name)
-            else if (element?.type === 'RestElement' && element.argument.type === 'Identifier')
-              identifiers.add(element.argument.name)
-            else if (element?.type === 'ObjectPattern' || element?.type === 'ArrayPattern')
-              handleVariableId(element)
-          }
-        }
+export function getIdentifierDeclarations(nodes: Statement[]) {
+  let result!: Set<string>
+  let programScopeUid: number
+  traverse(t.file(t.program(nodes)), {
+    Program(path) {
+      result = new Set(Object.keys(path.scope.bindings))
+      programScopeUid = (path.scope as any).uid
+    },
+    // FIXME: babel bug, temporary add TSEnumDeclaration and TSModuleDeclaration logic
+    TSEnumDeclaration(path) {
+      if ((path.scope as any).uid === programScopeUid)
+        result.add(path.node.id.name)
+    },
+    TSModuleDeclaration(path) {
+      if ((path.scope as any).uid === programScopeUid) {
+        const id = path.node.id
+        if (id.type === 'Identifier')
+          result.add(id.name)
       }
-
-      for (const declarator of node.declarations)
-        handleVariableId(declarator.id)
-    }
-    else if (node.type === 'FunctionDeclaration' || node.type === 'ClassDeclaration') {
-      if (node.id)
-        identifiers.add(node.id.name)
-    }
-    else if (node.type === 'TSEnumDeclaration') {
-      if (node.id)
-        identifiers.add(node.id.name)
-    }
-    // else {
-    //   console.log(node)
-    // }
-  }
-  return identifiers
+    },
+  })
+  return Array.from(result)
 }
 
+/**
+ * @deprecated use `getFileGlobals` instead
+ */
 export function getIdentifierUsages(node?: Expression | TSType | SpreadElement | PrivateName | Statement | null, identifiers = new Set<string>()) {
   if (!node)
     return identifiers
@@ -124,3 +106,27 @@ export function getIdentifierUsages(node?: Expression | TSType | SpreadElement |
   // }
   return identifiers
 }
+
+export function getFileGlobals(result: ParseResult<File>) {
+  let globals!: Set<string>
+  let programScopeUid: number
+  traverse(result, {
+    Program(path) {
+      globals = new Set(Object.keys((path.scope as any).globals))
+      programScopeUid = (path.scope as any).uid
+    },
+    // FIXME: babel bug, temporary add TSEnumDeclaration and TSModuleDeclaration logic
+    TSEnumDeclaration(path) {
+      if ((path.scope as any).uid === programScopeUid)
+        globals.delete(path.node.id.name)
+    },
+    TSModuleDeclaration(path) {
+      if ((path.scope as any).uid === programScopeUid) {
+        const id = path.node.id
+        if (id.type === 'Identifier')
+          globals.delete(id.name)
+      }
+    },
+  })
+  return Array.from(globals)
+}

src/core/parseSFC.ts

@@ -19,7 +19,7 @@ import type {
   ScriptSetupTransformOptions,
   ScriptTagMeta,
 } from '../types'
-import { getIdentifierUsages } from './identifiers'
+import { getFileGlobals } from './identifiers'
 import { parse } from './babel'
 import { exhaustiveCheckReturnUndefined, pascalize } from './utils'
 
@@ -137,12 +137,8 @@ function getDirectiveNames(node: TemplateChildNode): string[] {
 }
 
 function getFreeVariablesForText(input: string): string[] {
-  const identifiers = new Set<string>()
   const inputWithPrefix = input.trimStart()[0] === '{' ? `(${input})` : input
-
-  const nodes = parse(inputWithPrefix).program.body
-  nodes.forEach(node => getIdentifierUsages(node, identifiers))
-  return [...identifiers.values()]
+  return getFileGlobals(parse(inputWithPrefix))
 }
 
 function getFreeVariablesForPropsNode(

src/core/transformScriptSetup.ts

@@ -1,23 +1,31 @@
 import { capitalize } from '@vue/shared'
 import type { Node, ObjectExpression, Statement } from '@babel/types'
-import { partition } from '@antfu/utils'
+import { notNullish, partition, uniq } from '@antfu/utils'
 import type { ParsedSFC, ScriptSetupTransformOptions } from '../types'
 import { applyMacros } from './macros'
 import { getIdentifierDeclarations } from './identifiers'
 import { generate, t } from './babel'
-import { isNotNil, pascalize } from './utils'
+import { pascalize } from './utils'
 
-function isAsyncImport(node: any) {
-  if (node.type === 'VariableDeclaration') {
+function isAsyncImport(node: Statement) {
+  if (t.isVariableDeclaration(node)) {
     const declaration = node.declarations[0]
 
-    return declaration?.init?.callee?.name === 'defineAsyncComponent'
+    return (
+      declaration !== undefined
+      && t.isCallExpression(declaration.init)
+      && t.isIdentifier(declaration.init.callee)
+      && declaration.init.callee.name === 'defineAsyncComponent'
+    )
   }
 
   return false
 }
 
-export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransformOptions) {
+export function transformScriptSetup(
+  sfc: ParsedSFC,
+  options?: ScriptSetupTransformOptions,
+) {
   const { scriptSetup, script, template } = sfc
 
   const { nodes: body, props, expose } = applyMacros(scriptSetup.ast.body)
@@ -26,16 +34,17 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
     body,
     n =>
       isAsyncImport(n)
-     || n.type === 'ImportDeclaration'
-     || n.type === 'ExportNamedDeclaration'
-     || n.type.startsWith('TS'),
+      || t.isImportDeclaration(n)
+      || t.isExportNamedDeclaration(n)
+      || n.type.startsWith('TS'),
   )
 
   // get all identifiers in `<script setup>`
-  const declarations = new Set<string>()
-  getIdentifierDeclarations(hoisted, declarations)
-  getIdentifierDeclarations(setupBody, declarations)
-  const declarationArray = Array.from(declarations).filter(isNotNil)
+  const declarations = [
+    ...getIdentifierDeclarations(hoisted),
+    ...getIdentifierDeclarations(setupBody),
+  ]
+  const declarationArray = uniq(declarations).filter(notNullish)
 
   // filter out identifiers that are used in `<template>`
   const returns: ObjectExpression['properties'] = declarationArray
@@ -45,19 +54,24 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
       return t.objectProperty(id, id, false, true)
     })
 
-  const components = Array.from(template.components).map(component =>
-    declarationArray.find(declare => declare === component)
-    ?? declarationArray.find(declare => pascalize(declare) === component),
-  ).filter(isNotNil)
+  const components = Array.from(template.components)
+    .map(
+      component =>
+        declarationArray.find(declare => declare === component)
+        ?? declarationArray.find(declare => pascalize(declare) === component),
+    )
+    .filter(notNullish)
 
-  const directiveDeclaration = Array.from(template.directives).map((directive) => {
-    const identifier = declarationArray.find(declaration => declaration === `v${capitalize(directive)}`)
-    if (identifier === undefined)
-      return undefined
+  const directiveDeclaration = Array.from(template.directives)
+    .map((directive) => {
+      const identifier = declarationArray.find(
+        declaration => declaration === `v${capitalize(directive)}`,
+      )
+      if (identifier === undefined) return undefined
 
-    return { identifier, directive }
-  },
-  ).filter(isNotNil)
+      return { identifier, directive }
+    })
+    .filter(notNullish)
 
   // append `<script setup>` imports to `<script>`
 
@@ -71,10 +85,7 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
     if (node.type === 'ExportDefaultDeclaration') {
       hasBody = true
       return t.variableDeclaration('const', [
-        t.variableDeclarator(
-          __sfc,
-          node.declaration as any,
-        ),
+        t.variableDeclarator(__sfc, node.declaration as any),
       ])
     }
     return node
@@ -90,10 +101,7 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
   if (!hasBody) {
     ast.body.push(
       t.variableDeclaration('const', [
-        t.variableDeclarator(
-          __sfc,
-          t.objectExpression([]),
-        ),
+        t.variableDeclarator(__sfc, t.objectExpression([])),
       ]),
     )
   }
@@ -104,7 +112,8 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
     hasBody = true
     ast.body.push(
       t.expressionStatement(
-        t.assignmentExpression('=',
+        t.assignmentExpression(
+          '=',
           t.memberExpression(__sfc, t.identifier('props')),
           props as any,
         ),
@@ -126,15 +135,13 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
 
     ast.body.push(
       t.expressionStatement(
-        t.assignmentExpression('=',
+        t.assignmentExpression(
+          '=',
           t.memberExpression(__sfc, t.identifier('setup')),
-          t.arrowFunctionExpression([
-            t.identifier('__props'),
-            t.identifier('__ctx'),
-          ], t.blockStatement([
-            ...setupBody,
-            returnStatement as any,
-          ])),
+          t.arrowFunctionExpression(
+            [t.identifier('__props'), t.identifier('__ctx')],
+            t.blockStatement([...setupBody, returnStatement as any]),
+          ),
         ),
       ) as any,
     )
@@ -153,7 +160,8 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
 
     ast.body.push(
       t.expressionStatement(
-        t.assignmentExpression('=',
+        t.assignmentExpression(
+          '=',
           t.memberExpression(__sfc, t.identifier('components')),
           t.callExpression(
             t.memberExpression(t.identifier('Object'), t.identifier('assign')),
@@ -172,17 +180,20 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
   if (directiveDeclaration.length) {
     hasBody = true
     const directivesObject = t.objectExpression(
-      directiveDeclaration.map(({ directive, identifier }) => (t.objectProperty(
-        t.identifier(directive),
-        t.identifier(identifier),
-        false,
-        false,
-      ))),
+      directiveDeclaration.map(({ directive, identifier }) =>
+        t.objectProperty(
+          t.identifier(directive),
+          t.identifier(identifier),
+          false,
+          false,
+        ),
+      ),
     )
 
     ast.body.push(
       t.expressionStatement(
-        t.assignmentExpression('=',
+        t.assignmentExpression(
+          '=',
           t.memberExpression(__sfc, t.identifier('directives')),
           t.callExpression(
             t.memberExpression(t.identifier('Object'), t.identifier('assign')),
@@ -205,9 +216,7 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
 
   // re-export
   // `export default __sfc_main`
-  ast.body.push(
-    t.exportDefaultDeclaration(__sfc) as any,
-  )
+  ast.body.push(t.exportDefaultDeclaration(__sfc) as any)
 
   ast = options?.astTransforms?.post?.(ast, sfc) || ast
 

test/identifiers.test.ts

@@ -27,7 +27,7 @@ describe('identifiers', () => {
           sourceType: 'module',
         })
 
-        expect(getIdentifierDeclarations(ast.program.body)).toEqual(new Set(output))
+        expect(new Set(getIdentifierDeclarations(ast.program.body))).toEqual(new Set(output))
       })
     }
   })