From 837e4a5410e202733c3cbb121c87d4a7b40866aa Mon Sep 17 00:00:00 2001
From: ZHAO Jinxiang <xiaoxiangmoe@gmail.com>
Date: Sun, 28 Nov 2021 19:57:30 +0800
Subject: [PATCH] feat: add directive support (#75)

---
 src/core/parseSFC.ts                      | 36 ++++++++++++++-
 src/core/transformScriptSetup.ts          | 55 +++++++++++++++++++----
 src/core/utils.ts                         |  5 +++
 src/types.ts                              |  3 ++
 test/__snapshots__/transform.test.ts.snap | 37 +++++++++++++++
 test/fixtures/ComponentsDirectives.vue    | 16 +++++++
 6 files changed, 141 insertions(+), 11 deletions(-)
 create mode 100644 src/core/utils.ts
 create mode 100644 test/fixtures/ComponentsDirectives.vue

diff --git a/src/core/parseSFC.ts b/src/core/parseSFC.ts
index 4f0e8d2..ed539ce 100644
--- a/src/core/parseSFC.ts
+++ b/src/core/parseSFC.ts
@@ -1,15 +1,41 @@
 import { Parser as HTMLParser, ParserOptions as HTMLParserOptions } from 'htmlparser2'
 import type { ParserOptions } from '@babel/parser'
-import { camelize, capitalize, isHTMLTag, isSVGTag, isVoidTag } from '@vue/shared'
+import { camelize, isHTMLTag, isSVGTag, isVoidTag } from '@vue/shared'
 import { ParsedSFC, ScriptSetupTransformOptions, ScriptTagMeta } from '../types'
 import { getIdentifierUsages } from './identifiers'
 import { parse } from './babel'
+import { pascalize } from './utils'
 
 const multilineCommentsRE = /\/\*\s(.|[\r\n])*?\*\//gm
 const singlelineCommentsRE = /\/\/\s.*/g
 
+const BUILD_IN_DIRECTIVES = new Set([
+  'if',
+  'else',
+  'else-if',
+  'for',
+  'once',
+  'model',
+  'on',
+  'bind',
+  'slot',
+  'slot-scope',
+  'key',
+  'ref',
+  'text',
+  'html',
+  'show',
+  'pre',
+  'cloak',
+  // 'el',
+  // 'ref',
+])
+
 export function parseSFC(code: string, id?: string, options?: ScriptSetupTransformOptions): ParsedSFC {
+  /** foo-bar -> FooBar */
   const components = new Set<string>()
+  /** v-foo-bar -> fooBar */
+  const directives = new Set<string>()
   const expressions = new Set<string>()
   const identifiers = new Set<string>()
 
@@ -52,7 +78,7 @@ export function parseSFC(code: string, id?: string, options?: ScriptSetupTransfo
 
   function handleTemplateContent(name: string, attributes: Record<string, string>) {
     if (!isHTMLTag(name) && !isSVGTag(name) && !isVoidTag(name))
-      components.add(capitalize(camelize(name)))
+      components.add(pascalize((name)))
 
     Object.entries(attributes).forEach(([key, value]) => {
       if (!value)
@@ -64,6 +90,11 @@ export function parseSFC(code: string, id?: string, options?: ScriptSetupTransfo
         else
           expressions.add(`(${value})`)
       }
+      if (key.startsWith('v-')) {
+        const directiveName = key.slice('v-'.length).split(':')[0].split('.')[0]
+        if (!BUILD_IN_DIRECTIVES.has(directiveName))
+          directives.add(camelize(directiveName))
+      }
       if (key === 'ref')
         identifiers.add(value)
     })
@@ -183,6 +214,7 @@ export function parseSFC(code: string, id?: string, options?: ScriptSetupTransfo
     id,
     template: {
       components,
+      directives,
       identifiers,
     },
     scriptSetup,
diff --git a/src/core/transformScriptSetup.ts b/src/core/transformScriptSetup.ts
index 8f4833a..b9da054 100644
--- a/src/core/transformScriptSetup.ts
+++ b/src/core/transformScriptSetup.ts
@@ -1,4 +1,4 @@
-import { camelize, capitalize } from '@vue/shared'
+import { capitalize } from '@vue/shared'
 import type { Node, ObjectExpression, Statement } from '@babel/types'
 import generate from '@babel/generator'
 import { partition } from '@antfu/utils'
@@ -6,6 +6,7 @@ import { ParsedSFC, ScriptSetupTransformOptions } from '../types'
 import { applyMacros } from './macros'
 import { getIdentifierDeclarations } from './identifiers'
 import { t } from './babel'
+import { isNotNil, pascalize } from './utils'
 
 function isAsyncImport(node: any) {
   if (node.type === 'VariableDeclaration') {
@@ -35,22 +36,29 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
   const declarations = new Set<string>()
   getIdentifierDeclarations(hoisted, declarations)
   getIdentifierDeclarations(setupBody, declarations)
+  const declarationArray = Array.from(declarations).filter(isNotNil)
 
   // filter out identifiers that are used in `<template>`
-  const returns: ObjectExpression['properties'] = Array.from(declarations)
-    .filter(Boolean)
+  const returns: ObjectExpression['properties'] = declarationArray
     .filter(i => template.identifiers.has(i))
     .map((i) => {
       const id = t.identifier(i)
       return t.objectProperty(id, id, false, true)
     })
 
-  const components = Array.from(declarations)
-    .filter(Boolean)
-    .filter(i => template.components.has(i)
-      || template.components.has(camelize(i))
-      || template.components.has(capitalize(camelize(i))),
-    )
+  const components = Array.from(template.components).map(component =>
+    declarationArray.find(declare => declare === component)
+    ?? declarationArray.find(declare => pascalize(declare) === component),
+  ).filter(isNotNil)
+
+  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)
 
   // append `<script setup>` imports to `<script>`
 
@@ -160,6 +168,35 @@ export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransf
     )
   }
 
+  // inject directives
+  // `__sfc_main.directives = Object.assign({ ... }, __sfc_main.directives)`
+  if (directiveDeclaration.length) {
+    hasBody = true
+    const directivesObject = t.objectExpression(
+      directiveDeclaration.map(({ directive, identifier }) => (t.objectProperty(
+        t.identifier(directive),
+        t.identifier(identifier),
+        false,
+        false,
+      ))),
+    )
+
+    ast.body.push(
+      t.expressionStatement(
+        t.assignmentExpression('=',
+          t.memberExpression(__sfc, t.identifier('directives')),
+          t.callExpression(
+            t.memberExpression(t.identifier('Object'), t.identifier('assign')),
+            [
+              directivesObject,
+              t.memberExpression(__sfc, t.identifier('directives')),
+            ],
+          ),
+        ),
+      ) as any,
+    )
+  }
+
   if (!hasBody && !options?.astTransforms) {
     return {
       ast: null,
diff --git a/src/core/utils.ts b/src/core/utils.ts
new file mode 100644
index 0000000..8345bf1
--- /dev/null
+++ b/src/core/utils.ts
@@ -0,0 +1,5 @@
+import { camelize, capitalize } from '@vue/shared'
+
+export const pascalize = (str: string) => capitalize(camelize(str))
+
+export const isNotNil = <T>(value: T): value is NonNullable<T> => value != null
diff --git a/src/types.ts b/src/types.ts
index 05b3bf6..4f9edbb 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -17,7 +17,10 @@ export interface ScriptTagMeta {
 export interface ParsedSFC {
   id?: string
   template: {
+    /** foo-bar -> FooBar */
     components: Set<string>
+    /** v-foo-bar -> fooBar */
+    directives: Set<string>
     identifiers: Set<string>
   }
   scriptSetup: ScriptTagMeta
diff --git a/test/__snapshots__/transform.test.ts.snap b/test/__snapshots__/transform.test.ts.snap
index ccde47d..51e57a0 100644
--- a/test/__snapshots__/transform.test.ts.snap
+++ b/test/__snapshots__/transform.test.ts.snap
@@ -180,6 +180,43 @@ export default __sfc_main;
 "
 `;
 
+exports[`transform fixtures test/fixtures/ComponentsDirectives.vue 1`] = `
+"<template>
+  <div>
+    <FooView ref=\\"fooView\\" v-foo-bar=\\"a\\" v-demo:foo.a.b=\\"message\\"></FooView>
+    <router-view></router-view>
+  </div>
+</template>
+
+<script lang=\\"ts\\">
+import { ref } from '@vue/runtime-dom';
+import FooView from './FooView.vue';
+import { vFooBar, vDemo } from './directive';
+const __sfc_main = {};
+
+__sfc_main.setup = (__props, __ctx) => {
+  const fooView = ref<null | InstanceType<typeof FooView>>(null);
+  const a = ref(1);
+  const message = ref('hello');
+  return {
+    fooView,
+    a,
+    message
+  };
+};
+
+__sfc_main.components = Object.assign({
+  FooView
+}, __sfc_main.components);
+__sfc_main.directives = Object.assign({
+  fooBar: vFooBar,
+  demo: vDemo
+}, __sfc_main.directives);
+export default __sfc_main;
+</script>
+"
+`;
+
 exports[`transform fixtures test/fixtures/DynamicStyle.vue 1`] = `
 "<template>
   <div :style=\\"{ color, border: '1px' }\\" />
diff --git a/test/fixtures/ComponentsDirectives.vue b/test/fixtures/ComponentsDirectives.vue
new file mode 100644
index 0000000..5aa9c97
--- /dev/null
+++ b/test/fixtures/ComponentsDirectives.vue
@@ -0,0 +1,16 @@
+<template>
+  <div>
+    <FooView ref="fooView" v-foo-bar="a" v-demo:foo.a.b="message"></FooView>
+    <router-view></router-view>
+  </div>
+</template>
+
+<script setup lang="ts">
+import { ref } from '@vue/runtime-dom'
+import FooView from './FooView.vue'
+import { vFooBar, vDemo } from './directive'
+
+const fooView = ref<null | InstanceType<typeof FooView>>(null)
+const a = ref(1)
+const message = ref('hello')
+</script>