Support injected components inside MDX expressions (#398)

Closes #260

Author: remcohaszing

.changeset/large-dolls-obey.md

@@ -4,4 +4,4 @@
 "vscode-mdx": patch
 ---
 
-Support the `components` prop for MDX JSX tags.
+Support the `components` prop for MDX JSX tags and JSX expressions.

package.json

@@ -44,6 +44,7 @@
   "xo": {
     "prettier": true,
     "rules": {
+      "complexity": "off",
       "max-params": "off",
       "unicorn/prevent-abbreviations": "off"
     }

packages/language-service/lib/virtual-code.js

@@ -284,39 +284,6 @@ function processExports(mdx, node, mapping, esm) {
   return esm + '\n'
 }
 
-/**
- * @param {Program | undefined} expression
- * @returns {boolean}
- */
-function hasAwaitExpression(expression) {
-  let awaitExpression = false
-  if (expression) {
-    walk(expression, {
-      enter(node) {
-        if (
-          awaitExpression ||
-          node.type === 'ArrowFunctionExpression' ||
-          node.type === 'FunctionDeclaration' ||
-          node.type === 'FunctionExpression'
-        ) {
-          this.skip()
-          return
-        }
-
-        if (
-          node.type === 'AwaitExpression' ||
-          (node.type === 'ForOfStatement' && node.await)
-        ) {
-          awaitExpression = true
-          this.skip()
-        }
-      }
-    })
-  }
-
-  return awaitExpression
-}
-
 /**
  * @param {string} mdx
  * @param {Root} ast
@@ -472,6 +439,73 @@ function getEmbeddedCodes(mdx, ast, checkMdx, jsxImportSource) {
     updateMarkdownFromOffsets(startOffset, endOffset)
   }
 
+  /**
+   * @param {Program} program
+   * @param {number} lastIndex
+   * @returns {number}
+   */
+  function processJsxExpression(program, lastIndex) {
+    let newIndex = lastIndex
+    let functionNesting = 0
+    walk(program, {
+      enter(node) {
+        switch (node.type) {
+          case 'JSXIdentifier': {
+            if (!isInjectableComponent(node.name, variables)) {
+              return
+            }
+
+            jsx =
+              addOffset(jsxMapping, mdx, jsx, newIndex, node.start) +
+              '_components.'
+            newIndex = node.start
+            break
+          }
+
+          case 'ArrowFunctionExpression':
+          case 'FunctionDeclaration':
+          case 'FunctionExpression': {
+            functionNesting++
+            break
+          }
+
+          case 'AwaitExpression': {
+            if (!functionNesting) {
+              hasAwait = true
+            }
+
+            break
+          }
+
+          case 'ForOfStatement': {
+            if (!functionNesting) {
+              hasAwait ||= node.await
+            }
+
+            break
+          }
+
+          default:
+        }
+      },
+
+      leave(node) {
+        switch (node.type) {
+          case 'ArrowFunctionExpression':
+          case 'FunctionDeclaration':
+          case 'FunctionExpression': {
+            functionNesting--
+            break
+          }
+
+          default:
+        }
+      }
+    })
+
+    return newIndex
+  }
+
   visit(
     ast,
     (node) => {
@@ -525,38 +559,53 @@ function getEmbeddedCodes(mdx, ast, checkMdx, jsxImportSource) {
           }
 
           updateMarkdownFromOffsets(start, end)
+
+          let lastIndex = start + 1
+          jsx = addOffset(jsxMapping, mdx, jsx + jsxIndent, start, lastIndex)
           if (isInjectableComponent(node.name, variables)) {
-            const openingStart = start + 1
+            jsx += '_components.'
+          }
+
+          if (node.name) {
             jsx = addOffset(
               jsxMapping,
               mdx,
-              addOffset(jsxMapping, mdx, jsx + jsxIndent, start, openingStart) +
-                '_components.',
-              openingStart,
-              end
+              jsx,
+              lastIndex,
+              lastIndex + node.name.length
             )
-          } else {
-            jsx = addOffset(jsxMapping, mdx, jsx + jsxIndent, start, end)
+            lastIndex += node.name.length
+          }
+
+          for (const attribute of node.attributes) {
+            if (typeof attribute.value !== 'object') {
+              continue
+            }
+
+            const program = attribute.value?.data?.estree
+
+            if (program) {
+              lastIndex = processJsxExpression(program, lastIndex)
+            }
           }
 
+          jsx = addOffset(jsxMapping, mdx, jsx, lastIndex, end)
           break
         }
 
         case 'mdxFlowExpression':
         case 'mdxTextExpression': {
           updateMarkdownFromNode(node)
           const program = node.data?.estree
+          jsx += jsxIndent
 
-          if (program?.body.length === 0) {
-            jsx = addOffset(jsxMapping, mdx, jsx + jsxIndent, start, start + 1)
+          if (program?.body.length) {
+            const newIndex = processJsxExpression(program, start)
+            jsx = addOffset(jsxMapping, mdx, jsx, newIndex, end)
+          } else {
+            jsx = addOffset(jsxMapping, mdx, jsx, start, start + 1)
             jsx = addOffset(jsxMapping, mdx, jsx, end - 1, end)
             esm = addOffset(esmMapping, mdx, esm, start + 1, end - 1) + '\n'
-          } else {
-            if (program) {
-              hasAwait ||= hasAwaitExpression(program)
-            }
-
-            jsx = addOffset(jsxMapping, mdx, jsx + jsxIndent, start, end)
           }
 
           break