[refresh] Fix default export signatures for memo/forwardRef/hoc

packages/react-refresh/src/ReactFreshBabelPlugin.js

@@ -98,9 +98,20 @@ export default function (babel, opts = {}) {
             if (!foundInside) {
               return false;
             }
-            // const Foo = hoc1(hoc2(() => {}))
-            // export default memo(React.forwardRef(function() {}))
-            callback(inferredName, node, path);
+
+            // TODO: this is a hack, we should find a better way to detect this.
+            if (
+              firstArgPath.node.type === 'Identifier' &&
+              inferredName === '%default%'
+            ) {
+              // export default memo(function() {}))
+              // export default forwardRef(function() {}))
+              callback(innerName, node, path);
+            } else {
+              // const Foo = hoc1(hoc2(() => {}))
+              // export default memo(React.forwardRef(function() {}))
+              callback(inferredName, node, path);
+            }
             return true;
           }
           default: {

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

@@ -206,6 +206,44 @@
     ).toMatchSnapshot();
   });
 
+  it('registers memo default export', () => {
+    expect(
+      transform(`
+        function Component() {}
+        export default React.memo(Component);
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('registers forwardRef default export', () => {
+    expect(
+      transform(`
+        function Component() {}
+        export default React.forwardRef(Component);
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('registers memo default export', () => {
+    expect(
+      transform(`
+        import {memo} from 'react';
+        function Component() {}
+        export default memo(Component);
+    `),
+    ).toMatchSnapshot();
+  });
+
+  it('registers forwardRef default export', () => {
+    expect(
+      transform(`
+        import {forwardRef} from 'react';
+        function Component() {}
+        export default forwardRef(Component);
+    `),
+    ).toMatchSnapshot();
+  });
+
   it('registers likely HOCs with inline functions', () => {
     expect(
       transform(`

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

@@ -274,6 +274,123 @@ describe('ReactFreshIntegration', () => {
       }
     });
 
+    it('switches memo to forwardRef with different inner', async () => {
+      if (__DEV__) {
+        await render(`
+          const Child = () => {
+            return <h1>A1</h1>;
+          };
+
+          export default Child;
+        `);
+        let lastElement = container.firstChild;
+        expect(lastElement.textContent).toBe('A1');
+        await patch(`
+          const {memo} = React;
+
+          const Child2 = () => {
+            return <h1>A2</h1>;
+          };
+
+          export default memo(Child2);
+        `);
+        expect(container.firstChild !== lastElement).toBe(true);
+        expect(container.firstChild.textContent).toBe('A2');
+        lastElement = container.firstChild;
+
+        await patch(`
+          const {forwardRef} = React;
+
+          const Child3 = () => {
+            return <h1>A3</h1>;
+          };
+
+          export default forwardRef(Child3);
+        `);
+
+        expect(container.firstChild !== lastElement).toBe(true);
+        expect(container.firstChild.textContent).toBe('A3');
+        lastElement = container.firstChild;
+
+        await patch(`
+          const {memo} = React;
+
+          const Child4 = () => {
+            return <h1>A4</h1>;
+          };
+
+          export default memo(Child4);
+        `);
+
+        expect(container.firstChild !== lastElement).toBe(true);
+        expect(container.firstChild.textContent).toBe('A4');
+        lastElement = container.firstChild;
+
+        await patch(`
+          const Child5 = () => {
+            return <h1>A5</h1>;
+          };
+
+          export default Child5;
+        `);
+
+        expect(container.firstChild !== lastElement).toBe(true);
+        expect(container.firstChild.textContent).toBe('A5');
+      }
+    });
+
+    it('switches memo to forwardRef with same inner', async () => {
+      if (__DEV__) {
+        await render(`
+          const Child = () => {
+            return <h1>A1</h1>;
+          };
+
+          export default Child;
+        `);
+        let lastElement = container.firstChild;
+        expect(lastElement.textContent).toBe('A1');
+        await patch(`
+          const {memo} = React;
+
+          const Child = () => {
+            return <h1>A1</h1>;
+          };
+
+          export default memo(Child);
+        `);
+
+        expect(container.firstChild !== lastElement).toBe(true);
+        expect(container.firstChild.textContent).toBe('A1');
+        lastElement = container.firstChild;
+
+        await patch(`
+          const {forwardRef} = React;
+
+          const Child = () => {
+            return <h1>A1</h1>;
+          };
+
+          export default forwardRef(Child);
+        `);
+
+        expect(container.firstChild !== lastElement).toBe(true);
+        expect(container.firstChild.textContent).toBe('A1');
+        lastElement = container.firstChild;
+
+        await patch(`
+          const Child = () => {
+            return <h1>A1</h1>;
+          };
+
+          export default Child;
+        `);
+
+        expect(container.firstChild !== lastElement).toBe(true);
+        expect(container.firstChild.textContent).toBe('A1');
+      }
+    });
+
     it('reloads default export with named memo', async () => {
       if (__DEV__) {
         await render(`

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

@@ -268,11 +268,30 @@ const B = hoc(Foo);
 _c4 = B;
 var _c, _c2, _c3, _c4;
 $RefreshReg$(_c, "Foo");
-$RefreshReg$(_c2, "%default%");
+$RefreshReg$(_c2, "%default%$hoc");
 $RefreshReg$(_c3, "A");
 $RefreshReg$(_c4, "B");
 `;
 
+exports[`ReactFreshBabelPlugin registers forwardRef default export 1`] = `
+function Component() {}
+_c = Component;
+export default _c2 = React.forwardRef(Component);
+var _c, _c2;
+$RefreshReg$(_c, "Component");
+$RefreshReg$(_c2, "%default%$React.forwardRef");
+`;
+
+exports[`ReactFreshBabelPlugin registers forwardRef default export 2`] = `
+import { forwardRef } from 'react';
+function Component() {}
+_c = Component;
+export default _c2 = forwardRef(Component);
+var _c, _c2;
+$RefreshReg$(_c, "Component");
+$RefreshReg$(_c2, "%default%$forwardRef");
+`;
+
 exports[`ReactFreshBabelPlugin registers identifiers used in JSX at definition site 1`] = `
 import A from './A';
 import Store from './Store';
@@ -397,6 +416,25 @@ $RefreshReg$(_c2, "%default%$React.memo");
 $RefreshReg$(_c3, "%default%");
 `;
 
+exports[`ReactFreshBabelPlugin registers memo default export 1`] = `
+function Component() {}
+_c = Component;
+export default _c2 = React.memo(Component);
+var _c, _c2;
+$RefreshReg$(_c, "Component");
+$RefreshReg$(_c2, "%default%$React.memo");
+`;
+
+exports[`ReactFreshBabelPlugin registers memo default export 2`] = `
+import { memo } from 'react';
+function Component() {}
+_c = Component;
+export default _c2 = memo(Component);
+var _c, _c2;
+$RefreshReg$(_c, "Component");
+$RefreshReg$(_c2, "%default%$memo");
+`;
+
 exports[`ReactFreshBabelPlugin registers top-level exported function declarations 1`] = `
 export function Hello() {
   function handleClick() {}