Fix some performance by improving vdom diffing

Closes GH-2029. Closes GH-2062. Reviewed-by: Titus Wormer <[email protected]> Reviewed-by: Christian Murphy <[email protected]>
mdx-js · Jun 17, 2022 · 7bcd705 · 7bcd705 · vercel · Jun 17, 2022
1 parent ac44dc7
commit 7bcd705
Show file tree

Hide file tree

Showing 4 changed files with 133 additions and 94 deletions.
diff --git a/packages/mdx/lib/plugin/recma-document.js b/packages/mdx/lib/plugin/recma-document.js
@@ -231,7 +231,7 @@ export function recmaDocument(options = {}) {
           child.expression.type === 'JSXElement')
       ) {
         content = true
-        replacement.push(createMdxContent(child.expression))
+        replacement.push(...createMdxContent(child.expression, Boolean(layout)))
         // The following catch-all branch is because plugins might’ve added
         // other things.
         // Normally, we only have import/export/jsx, but just add whatever’s
@@ -244,7 +244,7 @@ export function recmaDocument(options = {}) {
 
     // If there was no JSX content at all, add an empty function.
     if (!content) {
-      replacement.push(createMdxContent())
+      replacement.push(...createMdxContent(undefined, Boolean(layout)))
     }
 
     exportedIdentifiers.push(['MDXContent', 'default'])
@@ -481,9 +481,10 @@ export function recmaDocument(options = {}) {
 
   /**
    * @param {Expression} [content]
-   * @returns {FunctionDeclaration}
+   * @param {boolean} [hasInternalLayout]
+   * @returns {FunctionDeclaration[]}
    */
-  function createMdxContent(content) {
+  function createMdxContent(content, hasInternalLayout) {
     /** @type {JSXElement} */
     const element = {
       type: 'JSXElement',
@@ -508,7 +509,12 @@ export function recmaDocument(options = {}) {
           openingElement: {
             type: 'JSXOpeningElement',
             name: {type: 'JSXIdentifier', name: '_createMdxContent'},
-            attributes: [],
+            attributes: [
+              {
+                type: 'JSXSpreadAttribute',
+                argument: {type: 'Identifier', name: 'props'}
+              }
+            ],
             selfClosing: true
           },
           closingElement: null,
@@ -518,7 +524,21 @@ export function recmaDocument(options = {}) {
     }
 
     // @ts-expect-error: JSXElements are expressions.
-    const consequent = /** @type {Expression} */ (element)
+    let result = /** @type {Expression} */ (element)
+
+    if (!hasInternalLayout) {
+      result = {
+        type: 'ConditionalExpression',
+        test: {type: 'Identifier', name: 'MDXLayout'},
+        consequent: result,
+        alternate: {
+          type: 'CallExpression',
+          callee: {type: 'Identifier', name: '_createMdxContent'},
+          arguments: [{type: 'Identifier', name: 'props'}],
+          optional: false
+        }
+      }
+    }
 
     let argument = content || {type: 'Literal', value: null}
 
@@ -535,44 +555,36 @@ export function recmaDocument(options = {}) {
       argument = argument.children[0]
     }
 
-    return {
-      type: 'FunctionDeclaration',
-      id: {type: 'Identifier', name: 'MDXContent'},
-      params: [
-        {
-          type: 'AssignmentPattern',
-          left: {type: 'Identifier', name: 'props'},
-          right: {type: 'ObjectExpression', properties: []}
+    return [
+      {
+        type: 'FunctionDeclaration',
+        id: {type: 'Identifier', name: '_createMdxContent'},
+        params: [{type: 'Identifier', name: 'props'}],
+        body: {
+          type: 'BlockStatement',
+          body: [{type: 'ReturnStatement', argument}]
         }
-      ],
-      body: {
-        type: 'BlockStatement',
-        body: [
-          {
-            type: 'ReturnStatement',
-            argument: {
-              type: 'ConditionalExpression',
-              test: {type: 'Identifier', name: 'MDXLayout'},
-              consequent,
-              alternate: {
-                type: 'CallExpression',
-                callee: {type: 'Identifier', name: '_createMdxContent'},
-                arguments: [],
-                optional: false
-              }
-            }
-          },
+      },
+      {
+        type: 'FunctionDeclaration',
+        id: {type: 'Identifier', name: 'MDXContent'},
+        params: [
           {
-            type: 'FunctionDeclaration',
-            id: {type: 'Identifier', name: '_createMdxContent'},
-            params: [],
-            body: {
-              type: 'BlockStatement',
-              body: [{type: 'ReturnStatement', argument}]
-            }
+            type: 'AssignmentPattern',
+            left: {type: 'Identifier', name: 'props'},
+            right: {type: 'ObjectExpression', properties: []}
           }
-        ]
+        ],
+        body: {
+          type: 'BlockStatement',
+          body: [
+            {
+              type: 'ReturnStatement',
+              argument: result
+            }
+          ]
+        }
       }
-    }
+    ]
   }
 }
diff --git a/packages/mdx/lib/plugin/recma-jsx-rewrite.js b/packages/mdx/lib/plugin/recma-jsx-rewrite.js
@@ -76,6 +76,10 @@ export function recmaJsxRewrite(options = {}) {
     walk(tree, {
       enter(_node) {
         const node = /** @type {Node} */ (_node)
+        const newScope = /** @type {Scope|undefined} */ (
+          // @ts-expect-error: periscopic doesn’t support JSX.
+          scopeInfo.map.get(node)
+        )
 
         if (
           node.type === 'FunctionDeclaration' ||
@@ -89,30 +93,26 @@ export function recmaJsxRewrite(options = {}) {
             references: {},
             node
           })
-        }
 
-        let fnScope = fnStack[0]
+          // MDXContent only ever contains MDXLayout
+          if (
+            isNamedFunction(node, 'MDXContent') &&
+            newScope &&
+            !inScope(newScope, 'MDXLayout')
+          ) {
+            fnStack[0].components.push('MDXLayout')
+          }
+        }
 
+        const fnScope = fnStack[0]
         if (
           !fnScope ||
-          (!isNamedFunction(fnScope.node, 'MDXContent') &&
+          (!isNamedFunction(fnScope.node, '_createMdxContent') &&
             !providerImportSource)
         ) {
           return
         }
 
-        if (
-          fnStack[1] &&
-          isNamedFunction(fnStack[1].node, '_createMdxContent')
-        ) {
-          fnScope = fnStack[1]
-        }
-
-        const newScope = /** @type {Scope|undefined} */ (
-          // @ts-expect-error: periscopic doesn’t support JSX.
-          scopeInfo.map.get(node)
-        )
-
         if (newScope) {
           newScope.node = node
           currentScope = newScope

diff --git a/packages/mdx/test/compile.js b/packages/mdx/test/compile.js
@@ -778,16 +778,16 @@ test('jsx', async () => {
     String(compileSync('*a*', {jsx: true})),
     [
       '/*@jsxRuntime automatic @jsxImportSource react*/',
+      'function _createMdxContent(props) {',
+      '  const _components = Object.assign({',
+      '    p: "p",',
+      '    em: "em"',
+      '  }, props.components);',
+      '  return <_components.p><_components.em>{"a"}</_components.em></_components.p>;',
+      '}',
       'function MDXContent(props = {}) {',
       '  const {wrapper: MDXLayout} = props.components || ({});',
-      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent /></MDXLayout> : _createMdxContent();',
-      '  function _createMdxContent() {',
-      '    const _components = Object.assign({',
-      '      p: "p",',
-      '      em: "em"',
-      '    }, props.components);',
-      '    return <_components.p><_components.em>{"a"}</_components.em></_components.p>;',
-      '  }',
+      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent {...props} /></MDXLayout> : _createMdxContent(props);',
       '}',
       'export default MDXContent;',
       ''
@@ -799,12 +799,12 @@ test('jsx', async () => {
     String(compileSync('<a {...b} c d="1" e={1} />', {jsx: true})),
     [
       '/*@jsxRuntime automatic @jsxImportSource react*/',
+      'function _createMdxContent(props) {',
+      '  return <a {...b} c d="1" e={1} />;',
+      '}',
       'function MDXContent(props = {}) {',
       '  const {wrapper: MDXLayout} = props.components || ({});',
-      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent /></MDXLayout> : _createMdxContent();',
-      '  function _createMdxContent() {',
-      '    return <a {...b} c d="1" e={1} />;',
-      '  }',
+      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent {...props} /></MDXLayout> : _createMdxContent(props);',
       '}',
       'export default MDXContent;',
       ''
@@ -816,15 +816,15 @@ test('jsx', async () => {
     String(compileSync('<><a:b /><c.d/></>', {jsx: true})),
     [
       '/*@jsxRuntime automatic @jsxImportSource react*/',
+      'function _createMdxContent(props) {',
+      '  const {c} = props.components || ({});',
+      '  if (!c) _missingMdxReference("c", false);',
+      '  if (!c.d) _missingMdxReference("c.d", true);',
+      '  return <><><a:b /><c.d /></></>;',
+      '}',
       'function MDXContent(props = {}) {',
       '  const {wrapper: MDXLayout} = props.components || ({});',
-      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent /></MDXLayout> : _createMdxContent();',
-      '  function _createMdxContent() {',
-      '    const {c} = props.components || ({});',
-      '    if (!c) _missingMdxReference("c", false);',
-      '    if (!c.d) _missingMdxReference("c.d", true);',
-      '    return <><><a:b /><c.d /></></>;',
-      '  }',
+      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent {...props} /></MDXLayout> : _createMdxContent(props);',
       '}',
       'export default MDXContent;',
       'function _missingMdxReference(id, component) {',
@@ -840,12 +840,12 @@ test('jsx', async () => {
     [
       '/*@jsxRuntime automatic @jsxImportSource react*/',
       '/*1*/',
+      'function _createMdxContent(props) {',
+      '  return <><>{"a "}{}{" b"}</></>;',
+      '}',
       'function MDXContent(props = {}) {',
       '  const {wrapper: MDXLayout} = props.components || ({});',
-      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent /></MDXLayout> : _createMdxContent();',
-      '  function _createMdxContent() {',
-      '    return <><>{"a "}{}{" b"}</></>;',
-      '  }',
+      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent {...props} /></MDXLayout> : _createMdxContent(props);',
       '}',
       'export default MDXContent;',
       ''
@@ -857,15 +857,15 @@ test('jsx', async () => {
     String(compileSync('{<a-b></a-b>}', {jsx: true})),
     [
       '/*@jsxRuntime automatic @jsxImportSource react*/',
+      'function _createMdxContent(props) {',
+      '  const _components = Object.assign({',
+      '    "a-b": "a-b"',
+      '  }, props.components);',
+      '  return <>{<_components.a-b></_components.a-b>}</>;',
+      '}',
       'function MDXContent(props = {}) {',
       '  const {wrapper: MDXLayout} = props.components || ({});',
-      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent /></MDXLayout> : _createMdxContent();',
-      '  function _createMdxContent() {',
-      '    const _components = Object.assign({',
-      '      "a-b": "a-b"',
-      '    }, props.components);',
-      '    return <>{<_components.a-b></_components.a-b>}</>;',
-      '  }',
+      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent {...props} /></MDXLayout> : _createMdxContent(props);',
       '}',
       'export default MDXContent;',
       ''
@@ -877,22 +877,49 @@ test('jsx', async () => {
     String(compileSync('Hello {props.name}', {jsx: true})),
     [
       '/*@jsxRuntime automatic @jsxImportSource react*/',
+      'function _createMdxContent(props) {',
+      '  const _components = Object.assign({',
+      '    p: "p"',
+      '  }, props.components);',
+      '  return <_components.p>{"Hello "}{props.name}</_components.p>;',
+      '}',
       'function MDXContent(props = {}) {',
       '  const {wrapper: MDXLayout} = props.components || ({});',
-      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent /></MDXLayout> : _createMdxContent();',
-      '  function _createMdxContent() {',
-      '    const _components = Object.assign({',
-      '      p: "p"',
-      '    }, props.components);',
-      '    return <_components.p>{"Hello "}{props.name}</_components.p>;',
-      '  }',
+      '  return MDXLayout ? <MDXLayout {...props}><_createMdxContent {...props} /></MDXLayout> : _createMdxContent(props);',
       '}',
       'export default MDXContent;',
       ''
     ].join('\n'),
     'should allow using props'
   )
 
+  assert.equal(
+    String(
+      compileSync(
+        'export default function Layout({components, ...props}) { return <section {...props} /> }\n\na',
+        {jsx: true}
+      )
+    ),
+    [
+      '/*@jsxRuntime automatic @jsxImportSource react*/',
+      'const MDXLayout = function Layout({components, ...props}) {',
+      '  return <section {...props} />;',
+      '};',
+      'function _createMdxContent(props) {',
+      '  const _components = Object.assign({',
+      '    p: "p"',
+      '  }, props.components);',
+      '  return <_components.p>{"a"}</_components.p>;',
+      '}',
+      'function MDXContent(props = {}) {',
+      '  return <MDXLayout {...props}><_createMdxContent {...props} /></MDXLayout>;',
+      '}',
+      'export default MDXContent;',
+      ''
+    ].join('\n'),
+    'should not have a conditional expression for MDXLayout when there is an internal layout'
+  )
+
   assert.match(
     String(compileSync("{<w x='y \" z' />}", {jsx: true})),
     /x="y &quot; z"/,

diff --git a/packages/rollup/test/index.test.js b/packages/rollup/test/index.test.js
@@ -34,7 +34,7 @@ test('@mdx-js/rollup', async () => {
 
   assert.equal(
     output[0].map ? output[0].map.mappings : undefined,
-    ';;;MAAaA,OAAU,GAAA,MAAAC,GAAA,CAAAC,QAAA,EAAA;AAAQ,EAAA,QAAA,EAAA,QAAA;;;;;;;;;;;;AAE7B,MAAA,QAAA,EAAA,CAAA,SAAA,EAAAD,GAAA,CAAA,OAAA,EAAA,EAAA,CAAA,CAAA;;;;;;;',
+    ';;;MAAaA,OAAU,GAAA,MAAAC,GAAA,CAAAC,QAAA,EAAA;AAAQ,EAAA,QAAA,EAAA,QAAA;;;;;;;AAE7B,IAAA,QAAA,EAAA,CAAA,SAAA,EAAAD,GAAA,CAAA,OAAA,EAAA,EAAA,CAAA,CAAA;;;;;;;;;;;;',
     'should add a source map'
   )