[Fresh] Babel plugin now handles HOCs + add integration tests (#15724)

* Refactor component search to prepare for deeper traversals

* Register HOCs with intermediate results

* Register components that are used as JSX types

* Add integration test skeleton

The integration test combines testing runtime together with the Babel plugin. It's a bit harder to debug because multiple things can go wrong, but it helps us build confidence that specific scenarios work well.

* Add HOC integration test and fix conflict with JSX transform

* Infer usage from createElement too

This helps us avoid dependency on the plugin order.

* Remove outdated comments

* Wrap tests in __DEV__

* Support export default hoc(...) for anonymous functions

* Fix test indentation

* Fix comment typo

* Use named function for test as this case is more important

packages/react-fresh/src/ReactFreshBabelPlugin.js

@@ -32,25 +32,164 @@ export default function(babel) {
     return typeof name === 'string' && name[0] >= 'A' && name[0] <= 'Z';
   }
 
-  function isComponentish(node) {
+  function findInnerComponents(inferredName, path, callback) {
+    const node = path.node;
     switch (node.type) {
-      case 'FunctionDeclaration':
-        return node.id !== null && isComponentishName(node.id.name);
-      case 'VariableDeclarator':
-        return (
-          isComponentishName(node.id.name) &&
-          node.init !== null &&
-          (node.init.type === 'FunctionExpression' ||
-            (node.init.type === 'ArrowFunctionExpression' &&
-              node.init.body.type !== 'ArrowFunctionExpression'))
+      case 'FunctionDeclaration': {
+        // function Foo() {}
+        // export function Foo() {}
+        // export default function Foo() {}
+        callback(inferredName, node.id, null);
+        return true;
+      }
+      case 'ArrowFunctionExpression': {
+        if (node.body.type === 'ArrowFunctionExpression') {
+          return false;
+        }
+        // let Foo = () => {}
+        // export default hoc1(hoc2(() => {}))
+        callback(inferredName, node, path);
+        return true;
+      }
+      case 'FunctionExpression': {
+        // let Foo = function() {}
+        // const Foo = hoc1(forwardRef(function renderFoo() {}))
+        // export default memo(function() {})
+        callback(inferredName, node, path);
+        return true;
+      }
+      case 'CallExpression': {
+        const argsPath = path.get('arguments');
+        if (argsPath === undefined || argsPath.length === 0) {
+          return false;
+        }
+        const calleePath = path.get('callee');
+        switch (calleePath.node.type) {
+          case 'MemberExpression':
+          case 'Identifier': {
+            const calleeSource = calleePath.getSource();
+            const firstArgPath = argsPath[0];
+            const innerName = inferredName + '$' + calleeSource;
+            const foundInside = findInnerComponents(
+              innerName,
+              firstArgPath,
+              callback,
+            );
+            if (!foundInside) {
+              return false;
+            }
+            // const Foo = hoc1(hoc2(() => {}))
+            // export default memo(React.forwardRef(function() {}))
+            callback(inferredName, node, path);
+            return true;
+          }
+          default: {
+            return false;
+          }
+        }
+      }
+      case 'VariableDeclarator': {
+        const init = node.init;
+        if (init === null) {
+          return false;
+        }
+        const name = node.id.name;
+        if (!isComponentishName(name)) {
+          return false;
+        }
+        if (init.type === 'Identifier' || init.type === 'MemberExpression') {
+          return false;
+        }
+        const initPath = path.get('init');
+        const foundInside = findInnerComponents(
+          inferredName,
+          initPath,
+          callback,
         );
-      default:
-        return false;
+        if (foundInside) {
+          return true;
+        }
+        // See if this identifier is used in JSX. Then it's a component.
+        const binding = path.scope.getBinding(name);
+        if (binding === undefined) {
+          return;
+        }
+        let isLikelyUsedAsType = false;
+        const referencePaths = binding.referencePaths;
+        for (let i = 0; i < referencePaths.length; i++) {
+          const ref = referencePaths[i];
+          if (
+            ref.node.type !== 'JSXIdentifier' &&
+            ref.node.type !== 'Identifier'
+          ) {
+            continue;
+          }
+          const refParent = ref.parent;
+          if (refParent.type === 'JSXOpeningElement') {
+            isLikelyUsedAsType = true;
+          } else if (refParent.type === 'CallExpression') {
+            const callee = refParent.callee;
+            let fnName;
+            switch (callee.type) {
+              case 'Identifier':
+                fnName = callee.name;
+                break;
+              case 'MemberExpression':
+                fnName = callee.property.name;
+                break;
+            }
+            switch (fnName) {
+              case 'createElement':
+              case 'jsx':
+              case 'jsxDEV':
+              case 'jsxs':
+                isLikelyUsedAsType = true;
+                break;
+            }
+          }
+          if (isLikelyUsedAsType) {
+            // const X = ... + later <X />
+            callback(inferredName, init, initPath);
+            return true;
+          }
+        }
+      }
     }
+    return false;
   }
 
   return {
     visitor: {
+      ExportDefaultDeclaration(path) {
+        const node = path.node;
+        const decl = node.declaration;
+        const declPath = path.get('declaration');
+        if (decl.type !== 'CallExpression') {
+          // For now, we only support possible HOC calls here.
+          // Named function declarations are handled in FunctionDeclaration.
+          // Anonymous direct exports like export default function() {}
+          // are currently ignored.
+          return;
+        }
+        // This code path handles nested cases like:
+        // export default memo(() => {})
+        // In those cases it is more plausible people will omit names
+        // so they're worth handling despite possible false positives.
+        // More importantly, it handles the named case:
+        // export default memo(function Named() {})
+        const inferredName = '%default%';
+        const programPath = path.parentPath;
+        findInnerComponents(
+          inferredName,
+          declPath,
+          (persistentID, targetExpr, targetPath) => {
+            const handle = createRegistration(programPath, persistentID);
+            targetPath.replaceWith(
+              t.assignmentExpression('=', handle, targetExpr),
+            );
+          },
+        );
+      },
       FunctionDeclaration(path) {
         let programPath;
         let insertAfterPath;
@@ -60,6 +199,9 @@ export default function(babel) {
             programPath = path.parentPath;
             break;
           case 'ExportNamedDeclaration':
+            insertAfterPath = path.parentPath;
+            programPath = insertAfterPath.parentPath;
+            break;
           case 'ExportDefaultDeclaration':
             insertAfterPath = path.parentPath;
             programPath = insertAfterPath.parentPath;
@@ -67,17 +209,25 @@ export default function(babel) {
           default:
             return;
         }
-        const maybeComponent = path.node;
-        if (!isComponentish(maybeComponent)) {
+        const id = path.node.id;
+        if (id === null) {
+          // We don't currently handle anonymous default exports.
           return;
         }
-        const functionName = path.node.id.name;
-        const handle = createRegistration(programPath, functionName);
-        insertAfterPath.insertAfter(
-          t.expressionStatement(
-            t.assignmentExpression('=', handle, path.node.id),
-          ),
-        );
+        const inferredName = id.name;
+        if (!isComponentishName(inferredName)) {
+          return;
+        }
+        // export function Named() {}
+        // function Named() {}
+        findInnerComponents(inferredName, path, (persistentID, targetExpr) => {
+          const handle = createRegistration(programPath, persistentID);
+          insertAfterPath.insertAfter(
+            t.expressionStatement(
+              t.assignmentExpression('=', handle, targetExpr),
+            ),
+          );
+        });
       },
       VariableDeclaration(path) {
         let programPath;
@@ -92,20 +242,21 @@ export default function(babel) {
           default:
             return;
         }
-        const declPath = path.get('declarations');
-        if (declPath.length !== 1) {
+        const declPaths = path.get('declarations');
+        if (declPaths.length !== 1) {
           return;
         }
-        const firstDeclPath = declPath[0];
-        const maybeComponent = firstDeclPath.node;
-        if (!isComponentish(maybeComponent)) {
-          return;
-        }
-        const functionName = maybeComponent.id.name;
-        const initPath = firstDeclPath.get('init');
-        const handle = createRegistration(programPath, functionName);
-        initPath.replaceWith(
-          t.assignmentExpression('=', handle, initPath.node),
+        const declPath = declPaths[0];
+        const inferredName = declPath.node.id.name;
+        findInnerComponents(
+          inferredName,
+          declPath,
+          (persistentID, targetExpr, targetPath) => {
+            const handle = createRegistration(programPath, persistentID);
+            targetPath.replaceWith(
+              t.assignmentExpression('=', handle, targetExpr),
+            );
+          },
         );
       },
       Program: {

packages/react-fresh/src/__tests__/ReactFresh-test.js

@@ -16,10 +16,10 @@ let ReactDOM;
 let ReactFreshRuntime;
 let Scheduler;
 let act;
-let lastRoot;
 
 describe('ReactFresh', () => {
   let container;
+  let lastRoot;
   let scheduleHotUpdate;
 
   beforeEach(() => {

packages/react-fresh/src/__tests__/ReactFreshBabelPlugin-test.js

@@ -186,4 +186,117 @@ describe('ReactFreshBabelPlugin', () => {
     `),
     ).toMatchSnapshot();
   });
+
+  it('registers likely HOCs with inline functions', () => {
+    expect(
+      transform(`
+        const A = forwardRef(function() {
+          return <h1>Foo</h1>;
+        });
+        const B = memo(React.forwardRef(() => {
+          return <h1>Foo</h1>;
+        }));
+        export default React.memo(forwardRef((props, ref) => {
+          return <h1>Foo</h1>;
+        }));
+    `),
+    ).toMatchSnapshot();
+    expect(
+      transform(`
+        export default React.memo(forwardRef(function (props, ref) {
+          return <h1>Foo</h1>;
+        }));
+    `),
+    ).toMatchSnapshot();
+    expect(
+      transform(`
+        export default React.memo(forwardRef(function Named(props, ref) {
+          return <h1>Foo</h1>;
+        }));
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('ignores higher-order functions that are not HOCs', () => {
+    expect(
+      transform(`
+        const throttledAlert = throttle(function() {
+          alert('Hi');
+        });
+        const TooComplex = (function() { return hello })(() => {});
+        if (cond) {
+          const Foo = thing(() => {});
+        }
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('registers identifiers used in JSX at definition site', () => {
+    // When in doubt, register variables that were used in JSX.
+    // Foo, Header, and B get registered.
+    // A doesn't get registered because it's not declared locally.
+    // Alias doesn't get registered because its definition is just an identifier.
+    expect(
+      transform(`
+        import A from './A';
+        import Store from './Store';
+
+        Store.subscribe();
+
+        const Header = styled.div\`color: red\`
+        const Factory = funny.factory\`\`;
+
+        let Alias1 = A;
+        let Alias2 = A.Foo;
+        const Dict = {};
+
+        function Foo() {
+          return (
+            <div><A /><B /><Alias1 /><Alias2 /><Header /><Dict.X /></div>
+          );
+        }
+
+        const B = hoc(A);
+        const NotAComponent = wow(A);
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('registers identifiers used in React.createElement at definition site', () => {
+    // When in doubt, register variables that were used in JSX.
+    // Foo, Header, and B get registered.
+    // A doesn't get registered because it's not declared locally.
+    // Alias doesn't get registered because its definition is just an identifier.
+    expect(
+      transform(`
+        import A from './A';
+        import Store from './Store';
+
+        Store.subscribe();
+
+        const Header = styled.div\`color: red\`
+        const Factory = funny.factory\`\`;
+
+        let Alias1 = A;
+        let Alias2 = A.Foo;
+        const Dict = {};
+
+        function Foo() {
+          return [
+            React.createElement(A),
+            React.createElement(B),
+            React.createElement(Alias1),
+            React.createElement(Alias2),
+            jsx(Header),
+            React.createElement(Dict.X),
+          ];
+        }
+
+        React.createContext(Store);
+
+        const B = hoc(A);
+        const NotAComponent = wow(A);
+    `),
+    ).toMatchSnapshot();
+  });
 });

packages/react-fresh/src/__tests__/ReactFreshIntegration-test.js

@@ -0,0 +1,247 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ */
+
+/* eslint-disable no-for-of-loops/no-for-of-loops */
+
+'use strict';
+
+let React;
+let ReactDOM;
+let ReactFreshRuntime;
+let act;
+
+let babel = require('babel-core');
+let freshPlugin = require('react-fresh/babel');
+
+describe('ReactFreshIntegration', () => {
+  let container;
+  let lastRoot;
+  let scheduleHotUpdate;
+
+  beforeEach(() => {
+    global.__REACT_DEVTOOLS_GLOBAL_HOOK__ = {
+      supportsFiber: true,
+      inject: injected => {
+        scheduleHotUpdate = injected.scheduleHotUpdate;
+      },
+      onCommitFiberRoot: (id, root) => {
+        lastRoot = root;
+      },
+      onCommitFiberUnmount: () => {},
+    };
+
+    jest.resetModules();
+    React = require('react');
+    ReactDOM = require('react-dom');
+    ReactFreshRuntime = require('react-fresh/runtime');
+    act = require('react-dom/test-utils').act;
+    container = document.createElement('div');
+    document.body.appendChild(container);
+  });
+
+  afterEach(() => {
+    document.body.removeChild(container);
+  });
+
+  function execute(source) {
+    const compiled = babel.transform(source, {
+      babelrc: false,
+      presets: ['react'],
+      plugins: [freshPlugin, 'transform-es2015-modules-commonjs'],
+    }).code;
+    const exportsObj = {};
+    // eslint-disable-next-line no-new-func
+    new Function('React', 'exports', '__register__', compiled)(
+      React,
+      exportsObj,
+      __register__,
+    );
+    return exportsObj.default;
+  }
+
+  function render(source) {
+    const Component = execute(source);
+    act(() => {
+      ReactDOM.render(<Component />, container);
+    });
+  }
+
+  function patch(source) {
+    execute(source);
+    const hotUpdate = ReactFreshRuntime.prepareUpdate();
+    scheduleHotUpdate(lastRoot, hotUpdate);
+  }
+
+  function __register__(type, id) {
+    ReactFreshRuntime.register(type, id);
+  }
+
+  it('reloads function declarations', () => {
+    if (__DEV__) {
+      render(`
+        function Parent() {
+          return <Child prop="A" />;
+        };
+
+        function Child({prop}) {
+          return <h1>{prop}1</h1>;
+        };
+
+        export default Parent;
+      `);
+      const el = container.firstChild;
+      expect(el.textContent).toBe('A1');
+      patch(`
+        function Parent() {
+          return <Child prop="B" />;
+        };
+
+        function Child({prop}) {
+          return <h1>{prop}2</h1>;
+        };
+
+        export default Parent;
+      `);
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('B2');
+    }
+  });
+
+  it('reloads arrow functions', () => {
+    if (__DEV__) {
+      render(`
+        const Parent = () => {
+          return <Child prop="A" />;
+        };
+
+        const Child = ({prop}) => {
+          return <h1>{prop}1</h1>;
+        };
+
+        export default Parent;
+      `);
+      const el = container.firstChild;
+      expect(el.textContent).toBe('A1');
+      patch(`
+        const Parent = () => {
+          return <Child prop="B" />;
+        };
+
+        const Child = ({prop}) => {
+          return <h1>{prop}2</h1>;
+        };
+
+        export default Parent;
+      `);
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('B2');
+    }
+  });
+
+  it('reloads a combination of memo and forwardRef', () => {
+    if (__DEV__) {
+      render(`
+        const {memo} = React;
+
+        const Parent = memo(React.forwardRef(function (props, ref) {
+          return <Child prop="A" ref={ref} />;
+        }));
+
+        const Child = React.memo(({prop}) => {
+          return <h1>{prop}1</h1>;
+        });
+
+        export default React.memo(Parent);
+      `);
+      const el = container.firstChild;
+      expect(el.textContent).toBe('A1');
+      patch(`
+        const {memo} = React;
+
+        const Parent = memo(React.forwardRef(function (props, ref) {
+          return <Child prop="B" ref={ref} />;
+        }));
+
+        const Child = React.memo(({prop}) => {
+          return <h1>{prop}2</h1>;
+        });
+
+        export default React.memo(Parent);
+      `);
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('B2');
+    }
+  });
+
+  it('reloads default export with named memo', () => {
+    if (__DEV__) {
+      render(`
+        const {memo} = React;
+
+        const Child = React.memo(({prop}) => {
+          return <h1>{prop}1</h1>;
+        });
+
+        export default memo(React.forwardRef(function Parent(props, ref) {
+          return <Child prop="A" ref={ref} />;
+        }));
+      `);
+      const el = container.firstChild;
+      expect(el.textContent).toBe('A1');
+      patch(`
+        const {memo} = React;
+
+        const Child = React.memo(({prop}) => {
+          return <h1>{prop}2</h1>;
+        });
+
+        export default memo(React.forwardRef(function Parent(props, ref) {
+          return <Child prop="B" ref={ref} />;
+        }));
+      `);
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('B2');
+    }
+  });
+
+  it('reloads HOCs if they return functions', () => {
+    if (__DEV__) {
+      render(`
+        function hoc(letter) {
+          return function() {
+            return <h1>{letter}1</h1>;
+          }
+        }
+
+        export default function Parent() {
+          return <Child />;
+        }
+
+        const Child = hoc('A');
+      `);
+      const el = container.firstChild;
+      expect(el.textContent).toBe('A1');
+      patch(`
+        function hoc(letter) {
+          return function() {
+            return <h1>{letter}2</h1>;
+          }
+        }
+
+        export default function Parent() {
+          return React.createElement(Child);
+        }
+
+        const Child = hoc('B');
+      `);
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('B2');
+    }
+  });
+});

packages/react-fresh/src/__tests__/__snapshots__/ReactFreshBabelPlugin-test.js.snap

@@ -33,6 +33,19 @@ let D = bar && (() => {
 });"
 `;
 
+exports[`ReactFreshBabelPlugin ignores higher-order functions that are not HOCs 1`] = `
+"
+const throttledAlert = throttle(function () {
+  alert('Hi');
+});
+const TooComplex = function () {
+  return hello;
+}(() => {});
+if (cond) {
+  const Foo = thing(() => {});
+}"
+`;
+
 exports[`ReactFreshBabelPlugin ignores unnamed function declarations 1`] = `
 "
 export default function () {}"
@@ -45,6 +58,131 @@ function hello() {
 }"
 `;
 
+exports[`ReactFreshBabelPlugin registers identifiers used in JSX at definition site 1`] = `
+"
+import A from './A';
+import Store from './Store';
+
+Store.subscribe();
+
+const Header = _c = styled.div\`color: red\`;
+const Factory = funny.factory\`\`;
+
+let Alias1 = A;
+let Alias2 = A.Foo;
+const Dict = {};
+
+function Foo() {
+  return <div><A /><B /><Alias1 /><Alias2 /><Header /><Dict.X /></div>;
+}
+
+_c2 = Foo;
+const B = _c3 = hoc(A);
+const NotAComponent = wow(A);
+
+var _c, _c2, _c3;
+
+__register__(_c, 'Header');
+
+__register__(_c2, 'Foo');
+
+__register__(_c3, 'B');"
+`;
+
+exports[`ReactFreshBabelPlugin registers identifiers used in React.createElement at definition site 1`] = `
+"
+import A from './A';
+import Store from './Store';
+
+Store.subscribe();
+
+const Header = _c = styled.div\`color: red\`;
+const Factory = funny.factory\`\`;
+
+let Alias1 = A;
+let Alias2 = A.Foo;
+const Dict = {};
+
+function Foo() {
+  return [React.createElement(A), React.createElement(B), React.createElement(Alias1), React.createElement(Alias2), jsx(Header), React.createElement(Dict.X)];
+}
+
+_c2 = Foo;
+React.createContext(Store);
+
+const B = _c3 = hoc(A);
+const NotAComponent = wow(A);
+
+var _c, _c2, _c3;
+
+__register__(_c, 'Header');
+
+__register__(_c2, 'Foo');
+
+__register__(_c3, 'B');"
+`;
+
+exports[`ReactFreshBabelPlugin registers likely HOCs with inline functions 1`] = `
+"
+const A = _c2 = forwardRef(_c = function () {
+  return <h1>Foo</h1>;
+});
+const B = _c5 = memo(_c4 = React.forwardRef(_c3 = () => {
+  return <h1>Foo</h1>;
+}));
+export default _c8 = React.memo(_c7 = forwardRef(_c6 = (props, ref) => {
+  return <h1>Foo</h1>;
+}));
+
+var _c, _c2, _c3, _c4, _c5, _c6, _c7, _c8;
+
+__register__(_c, \\"A$forwardRef\\");
+
+__register__(_c2, \\"A\\");
+
+__register__(_c3, \\"B$memo$React.forwardRef\\");
+
+__register__(_c4, \\"B$memo\\");
+
+__register__(_c5, \\"B\\");
+
+__register__(_c6, \\"%default%$React.memo$forwardRef\\");
+
+__register__(_c7, \\"%default%$React.memo\\");
+
+__register__(_c8, \\"%default%\\");"
+`;
+
+exports[`ReactFreshBabelPlugin registers likely HOCs with inline functions 2`] = `
+"
+export default _c3 = React.memo(_c2 = forwardRef(_c = function (props, ref) {
+  return <h1>Foo</h1>;
+}));
+
+var _c, _c2, _c3;
+
+__register__(_c, \\"%default%$React.memo$forwardRef\\");
+
+__register__(_c2, \\"%default%$React.memo\\");
+
+__register__(_c3, \\"%default%\\");"
+`;
+
+exports[`ReactFreshBabelPlugin registers likely HOCs with inline functions 3`] = `
+"
+export default _c3 = React.memo(_c2 = forwardRef(_c = function Named(props, ref) {
+  return <h1>Foo</h1>;
+}));
+
+var _c, _c2, _c3;
+
+__register__(_c, \\"%default%$React.memo$forwardRef\\");
+
+__register__(_c2, \\"%default%$React.memo\\");
+
+__register__(_c3, \\"%default%\\");"
+`;
+
 exports[`ReactFreshBabelPlugin registers top-level exported function declarations 1`] = `
 "
 export function Hello() {