feat(babel-sugar-inject-h): add support for functional jsx h injections

packages/babel-sugar-inject-h/src/index.js

@@ -1,35 +1,61 @@
 import syntaxJsx from '@babel/plugin-syntax-jsx'
 
+const PLUGIN_DATA_PREFIX = `@vue/babel-sugar-inject-h_${Date.now()}`
+
+// helpers for ast custom data get/set
+const getv = (p, k) => p.getData(`${PLUGIN_DATA_PREFIX}_${k}`)
+const setv = (p, k, v) => p.setData(`${PLUGIN_DATA_PREFIX}_${k}`, v)
+
 /**
- * Check if first parameter is `h`
+ * Get the index of the parameter `h`
  * @param t
- * @param path ObjectMethod | ClassMethod
- * @returns boolean
+ * @param path ObjectMethod | ClassMethod | Function
+ * @returns number -1 if not found
  */
-const firstParamIsH = (t, path) => {
+const indexOfParamH = (t, path) => {
   const params = path.get('params')
-  return params.length && t.isIdentifier(params[0]) && params[0].node.name === 'h'
+  return params.length ? params.findIndex(p => t.isIdentifier(p) && p.node.name === 'h') : -1
 }
 
 /**
- * Check if body contains JSX
- * @param t
- * @param path ObjectMethod | ClassMethod
- * @returns boolean
+ * Check if expression is an object member function
  */
-const hasJSX = (t, path) => {
-  const JSXChecker = {
-    hasJSX: false,
+const isObjectMemberFunc = (t, path) => t.isFunction(path) && t.isObjectMember(path.parentPath)
+
+/**
+ * Check if expression is an object function-typed member
+ */
+const isMemberFunction = (t, path) => t.isObjectMethod(path) || t.isClassMethod(path) || isObjectMemberFunc(t, path)
+
+// find JSX, returns the first walked jsx expression
+const findJSXElement = (t, path) => {
+  let elem = null
+  path.traverse({
+    JSXElement (p) {
+      elem = p
+      p.stop()
+    }
+  })
+  return elem
+}
+
+/**
+ * Get function-typed ancestry of the specific node range
+ * @param path JSXElement
+ * @param root The last boundary node
+ * @returns the ancestry paths
+ */
+const getFuncAncestry = (t, path, root) => {
+  const paths = []
+  if (path !== root) {
+    while (path = path.parentPath) {
+      if (t.isFunction(path)) {
+        paths.push(path)
+      }
+      if (path === root) break
+    }
   }
-  path.traverse(
-    {
-      JSXElement() {
-        this.hasJSX = true
-      },
-    },
-    JSXChecker,
-  )
-  return JSXChecker.hasJSX
+  return paths
 }
 
 /**
@@ -38,14 +64,56 @@ const hasJSX = (t, path) => {
  * @param path ObjectMethod | ClassMethod
  * @returns boolean
  */
-const isInsideJSXExpression = (t, path) => {
-  if (!path.parentPath) {
-    return false
-  }
-  if (t.isJSXExpressionContainer(path.parentPath)) {
-    return true
+const isInsideJSXExpression = (t, path) => path.findParent(p => p && t.isJSX(p)) !== null
+
+/**
+ * Cleanup and reduce stack that `distance` gt than reference's
+ * @param stack
+ * @param reference The reference node to match
+ */
+const cleanupStack = (stack, path) => {
+  const ref = getv(path, 'distance')
+  stack.forEach(p => {
+    if (getv(p, 'distance') > ref) setv(p, 'hasJSX', false)
+  })
+}
+
+/**
+ * Prepend parameter `h` to function params (ensure as the first parameter)
+ *
+ * @param path {ArrowFunctionExpression|FunctionExpression}
+ */
+const addParamH = (t, path) => {
+  path.node.params = [t.identifier('h')].concat(path.node.params)
+}
+
+/**
+ * Prepend `h` variable to function body, i.e. const h = xxx;
+ *
+ * @param path {ObjectMethod|ClassMethod|FunctionExpression}
+ */
+const patchVariableH = (t, path) => {
+  const funcName = path.isFunctionExpression()
+    ? path.parent.key.name
+    : path.node.key.name
+  const isRender = funcName === 'render'
+
+  if (isRender) {
+    addParamH(t, path)
+    return
   }
-  return isInsideJSXExpression(t, path.parentPath)
+
+  path
+    .get('body')
+    .unshiftContainer(
+      'body',
+      t.variableDeclaration('const', [
+        t.variableDeclarator(
+          t.identifier('h'),
+          t.memberExpression(t.thisExpression(), t.identifier('$createElement')),
+        ),
+      ]),
+    )
 }
 
 export default babel => {
@@ -54,28 +122,67 @@ export default babel => {
   return {
     inherits: syntaxJsx,
     visitor: {
-      Program(path1) {
-        path1.traverse({
-          'ObjectMethod|ClassMethod'(path) {
-            if (firstParamIsH(t, path) || !hasJSX(t, path) || isInsideJSXExpression(t, path)) {
-              return
-            }
+      Program (p) {
+        const stack = []
+        p.traverse({
+          'ObjectMethod|ClassMethod|FunctionDeclaration|FunctionExpression|ArrowFunctionExpression': {
+            enter (path) {
+              const jsxElem = findJSXElement(t, path)
+              if (!jsxElem || isInsideJSXExpression(t, path)) {
+                return
+              }
 
-            const isRender = path.node.key.name === 'render'
-
-            path
-              .get('body')
-              .unshiftContainer(
-                'body',
-                t.variableDeclaration('const', [
-                  t.variableDeclarator(
-                    t.identifier('h'),
-                    isRender
-                      ? t.memberExpression(t.identifier('arguments'), t.numericLiteral(0), true)
-                      : t.memberExpression(t.thisExpression(), t.identifier('$createElement')),
-                  ),
-                ]),
-              )
+              const ancestry = getFuncAncestry(t, jsxElem, path)
+              const isObjFn = isMemberFunction(t, path)
+
+              // check if JSX expression is in method
+              if (!isObjFn && ancestry.some(p => isMemberFunction(t, p))) {
+                return
+              }
+
+              // add to pending stack
+              stack.push(path)
+
+              setv(path, 'hasJSX', true)
+              setv(path, 'distance', ancestry.length)
+
+              // check params already has param `h`
+              if (indexOfParamH(t, path) !== -1) {
+                setv(path, 'fixed', true)
+                return
+              }
+
+              if (isObjFn) {
+                if (path.isArrowFunctionExpression()) {
+                  addParamH(t, path)
+                } else {
+                  patchVariableH(t, path)
+                }
+                setv(path, 'fixed', true)
+              }
+            },
+            exit (path) {
+              if (!getv(path, 'hasJSX')) {
+                return
+              }
+              stack.pop()
+
+              // skip and cancel remaining nodes if `h` has fixed
+              if (getv(path, 'fixed')) {
+                cleanupStack(stack, path)
+                return
+              }
+
+              // skip, functional JSX `h` should to be fixed in top of pending stack
+              if (!isObjectMemberFunc(t, path) && stack.some(p => getv(p, 'hasJSX'))) {
+                return
+              }
+
+              addParamH(t, path)
+              setv(path, 'fixed', true)
+
+              cleanupStack(stack, path)
+            }
           }
         })
       }