Support MDX layout (#365)

Closes #362

.changeset/purple-mice-rhyme.md

@@ -0,0 +1,5 @@
+---
+"@mdx-js/language-service": minor
+---
+
+Support MDXLayout

package.json

@@ -42,15 +42,6 @@
     ]
   },
   "xo": {
-    "prettier": true,
-    "env": [
-      "browser",
-      "es2021",
-      "node"
-    ],
-    "rules": {
-      "@typescript-eslint/prefer-nullish-coalescing": "off",
-      "max-params": "off"
-    }
+    "prettier": true
   }
 }

packages/language-service/lib/language-module.js

@@ -3,8 +3,10 @@
  * @typedef {import('@volar/language-core').LanguagePlugin} LanguagePlugin
  * @typedef {import('@volar/language-core').Mapping<CodeInformation>} Mapping
  * @typedef {import('@volar/language-core').VirtualFile} VirtualFile
+ * @typedef {import('estree').ExportDefaultDeclaration} ExportDefaultDeclaration
  * @typedef {import('mdast').Root} Root
  * @typedef {import('mdast').RootContent} RootContent
+ * @typedef {import('mdast-util-mdxjs-esm').MdxjsEsm} MdxjsEsm
  * @typedef {import('typescript').IScriptSnapshot} IScriptSnapshot
  * @typedef {import('unified').PluggableList} PluggableList
  * @typedef {import('unified').Processor<Root>} Processor
@@ -16,6 +18,24 @@ import remarkMdx from 'remark-mdx'
 import remarkParse from 'remark-parse'
 import {unified} from 'unified'
 
+/**
+ * @param {string} propsName
+ */
+const layoutJsDoc = (propsName) => `
+/** @typedef {MDXContentProps & { children: JSX.Element }} MDXLayoutProps */
+
+/**
+ * There is one special component: [MDX layout](https://mdxjs.com/docs/using-mdx/#layout).
+ * If it is defined, it’s used to wrap all content.
+ * A layout can be defined from within MDX using a default export.
+ *
+ * @param {{readonly [K in keyof MDXLayoutProps]: MDXLayoutProps[K]}} ${propsName}
+ *   The [props](https://mdxjs.com/docs/using-mdx/#props) that have been passed to the MDX component.
+ *   In addition, the MDX layout receives the \`children\` prop, which contains the rendered MDX content.
+ * @returns {JSX.Element}
+ *   The MDX content wrapped in the layout.
+ */`
+
 const componentStart = `
 /**
  * Render the MDX contents.
@@ -67,6 +87,124 @@ function addOffset(mapping, sourceOffset, generatedOffset, length) {
   mapping.lengths.push(length)
 }
 
+/**
+ * @param {string} haystack
+ * @param {string} needle
+ * @param {number} start
+ */
+function findIndexAfter(haystack, needle, start) {
+  for (let index = start; index < haystack.length; index++) {
+    if (haystack[index] === needle) {
+      return index
+    }
+  }
+
+  return -1
+}
+
+/**
+ * @param {ExportDefaultDeclaration} node
+ */
+function getPropsName(node) {
+  const {declaration} = node
+  const {type} = declaration
+
+  if (
+    type !== 'ArrowFunctionExpression' &&
+    type !== 'FunctionDeclaration' &&
+    type !== 'FunctionExpression'
+  ) {
+    return
+  }
+
+  if (declaration.params.length === 1) {
+    const parameter = declaration.params[0]
+    if (parameter.type === 'Identifier') {
+      return parameter.name
+    }
+  }
+
+  return 'props'
+}
+
+/**
+ * Process exports of an MDX ESM node.
+ *
+ * @param {string} mdx
+ *   The full MDX code to process.
+ * @param {MdxjsEsm} node
+ *   The MDX ESM node to process.
+ * @param {Mapping} mapping
+ *   The Volar mapping to add offsets to.
+ * @param {string} esm
+ *   The virtual ESM code up to the point this function was called.
+ * @returns {string}
+ *   The updated virtual ESM code.
+ */
+function processExports(mdx, node, mapping, esm) {
+  const start = node.position?.start?.offset
+  const end = node.position?.end?.offset
+
+  if (start === undefined || end === undefined) {
+    return esm
+  }
+
+  const body = node.data?.estree?.body
+
+  if (!body?.length) {
+    addOffset(mapping, start, esm.length, end - start)
+    return esm + mdx.slice(start, end) + '\n'
+  }
+
+  for (const child of body) {
+    if (child.type === 'ExportDefaultDeclaration') {
+      const propsName = getPropsName(child)
+      if (propsName) {
+        esm += layoutJsDoc(propsName)
+      }
+
+      esm += '\nconst MDXLayout = '
+      addOffset(
+        mapping,
+        child.declaration.start,
+        esm.length,
+        child.end - child.declaration.start
+      )
+      esm += mdx.slice(child.declaration.start, child.end) + '\n'
+      continue
+    }
+
+    if (child.type === 'ExportNamedDeclaration' && child.source) {
+      const {specifiers} = child
+      for (let index = 0; index < specifiers.length; index++) {
+        const specifier = specifiers[index]
+        if (specifier.local.name === 'default') {
+          addOffset(mapping, start, esm.length, specifier.start - start)
+          esm += mdx.slice(start, specifier.start)
+          const nextPosition =
+            index === specifiers.length - 1
+              ? specifier.end
+              : findIndexAfter(mdx, ',', specifier.end) + 1
+          addOffset(mapping, nextPosition, esm.length, end - nextPosition)
+          return (
+            esm +
+            mdx.slice(nextPosition, end) +
+            '\nimport {' +
+            specifier.exported.name +
+            ' as MDXLayout} from ' +
+            JSON.stringify(child.source.value)
+          )
+        }
+      }
+    }
+
+    addOffset(mapping, child.start, esm.length, child.end - child.start)
+    esm += mdx.slice(child.start, child.end) + '\n'
+  }
+
+  return esm
+}
+
 /**
  * @param {string} fileName
  * @param {IScriptSnapshot} snapshot
@@ -250,8 +388,7 @@ function getVirtualFiles(fileName, snapshot, ts, processor) {
 
         case 'mdxjsEsm': {
           updateMarkdownFromNode(node)
-          addOffset(esmMapping, start, esm.length, end - start)
-          esm += mdx.slice(start, end) + '\n'
+          esm = processExports(mdx, node, esmMapping, esm)
           break
         }