[Fast Refresh] Support callthrough HOCs (#21104)

* [Fast Refresh] Support callthrough HOCs

* Add a newly failing testing to demonstrate the flaw

This shows why my initial approach doesn't make sense.

* Attach signatures at every nesting level

* Sign nested memo/forwardRef too

* Add an IIFE test

This is not a case that is important for Fast Refresh, but we shouldn't change the code semantics. This case shows the transform isn't quite correct. It's wrapping the call at the wrong place.

* Find HOCs above more precisely

This fixes a false positive that was causing an IIFE to be wrapped in the wrong place, which made the wrapping unsafe.

* Be defensive against non-components being passed to setSignature

* Fix lint

packages/react-refresh/src/ReactFreshBabelPlugin.js

@@ -333,6 +333,38 @@ export default function(babel, opts = {}) {
     return args;
   }
 
+  function findHOCCallPathsAbove(path) {
+    const calls = [];
+    while (true) {
+      if (!path) {
+        return calls;
+      }
+      const parentPath = path.parentPath;
+      if (!parentPath) {
+        return calls;
+      }
+      if (
+        // hoc(_c = function() { })
+        parentPath.node.type === 'AssignmentExpression' &&
+        path.node === parentPath.node.right
+      ) {
+        // Ignore registrations.
+        path = parentPath;
+        continue;
+      }
+      if (
+        // hoc1(hoc2(...))
+        parentPath.node.type === 'CallExpression' &&
+        path.node !== parentPath.node.callee
+      ) {
+        calls.push(parentPath);
+        path = parentPath;
+        continue;
+      }
+      return calls; // Stop at other types.
+    }
+  }
+
   const seenForRegistration = new WeakSet();
   const seenForSignature = new WeakSet();
   const seenForOutro = new WeakSet();
@@ -630,13 +662,16 @@ export default function(babel, opts = {}) {
             // Result: let Foo = () => {}; __signature(Foo, ...);
           } else {
             // let Foo = hoc(() => {})
-            path.replaceWith(
-              t.callExpression(
-                sigCallID,
-                createArgumentsForSignature(node, signature, path.scope),
-              ),
-            );
-            // Result: let Foo = hoc(__signature(() => {}, ...))
+            const paths = [path, ...findHOCCallPathsAbove(path)];
+            paths.forEach(p => {
+              p.replaceWith(
+                t.callExpression(
+                  sigCallID,
+                  createArgumentsForSignature(p.node, signature, p.scope),
+                ),
+              );
+            });
+            // Result: let Foo = __signature(hoc(__signature(() => {}, ...)), ...)
           }
         },
       },

packages/react-refresh/src/ReactFreshRuntime.js

@@ -355,12 +355,25 @@ export function setSignature(
   getCustomHooks?: () => Array<Function>,
 ): void {
   if (__DEV__) {
-    allSignaturesByType.set(type, {
-      forceReset,
-      ownKey: key,
-      fullKey: null,
-      getCustomHooks: getCustomHooks || (() => []),
-    });
+    if (!allSignaturesByType.has(type)) {
+      allSignaturesByType.set(type, {
+        forceReset,
+        ownKey: key,
+        fullKey: null,
+        getCustomHooks: getCustomHooks || (() => []),
+      });
+    }
+    // Visit inner types because we might not have signed them.
+    if (typeof type === 'object' && type !== null) {
+      switch (getProperty(type, '$$typeof')) {
+        case REACT_FORWARD_REF_TYPE:
+          setSignature(type.render, key, forceReset, getCustomHooks);
+          break;
+        case REACT_MEMO_TYPE:
+          setSignature(type.type, key, forceReset, getCustomHooks);
+          break;
+      }
+    }
   } else {
     throw new Error(
       'Unexpected call to React Refresh in a production environment.',
@@ -609,57 +622,58 @@ export function _getMountedRootCount() {
 // function Hello() {
 //   const [foo, setFoo] = useState(0);
 //   const value = useCustomHook();
-//   _s(); /* Second call triggers collecting the custom Hook list.
+//   _s(); /* Call without arguments triggers collecting the custom Hook list.
 //          * This doesn't happen during the module evaluation because we
 //          * don't want to change the module order with inline requires.
 //          * Next calls are noops. */
 //   return <h1>Hi</h1>;
 // }
 //
-// /* First call specifies the signature: */
+// /* Call with arguments attaches the signature to the type: */
 // _s(
 //   Hello,
 //   'useState{[foo, setFoo]}(0)',
 //   () => [useCustomHook], /* Lazy to avoid triggering inline requires */
 // );
-type SignatureStatus = 'needsSignature' | 'needsCustomHooks' | 'resolved';
 export function createSignatureFunctionForTransform() {
   if (__DEV__) {
-    // We'll fill in the signature in two steps.
-    // First, we'll know the signature itself. This happens outside the component.
-    // Then, we'll know the references to custom Hooks. This happens inside the component.
-    // After that, the returned function will be a fast path no-op.
-    let status: SignatureStatus = 'needsSignature';
     let savedType;
     let hasCustomHooks;
+    let didCollectHooks = false;
     return function<T>(
       type: T,
       key: string,
       forceReset?: boolean,
       getCustomHooks?: () => Array<Function>,
-    ): T {
-      switch (status) {
-        case 'needsSignature':
-          if (type !== undefined) {
-            // If we received an argument, this is the initial registration call.
-            savedType = type;
-            hasCustomHooks = typeof getCustomHooks === 'function';
-            setSignature(type, key, forceReset, getCustomHooks);
-            // The next call we expect is from inside a function, to fill in the custom Hooks.
-            status = 'needsCustomHooks';
-          }
-          break;
-        case 'needsCustomHooks':
-          if (hasCustomHooks) {
-            collectCustomHooksForSignature(savedType);
-          }
-          status = 'resolved';
-          break;
-        case 'resolved':
-          // Do nothing. Fast path for all future renders.
-          break;
+    ): T | void {
+      if (typeof key === 'string') {
+        // We're in the initial phase that associates signatures
+        // with the functions. Note this may be called multiple times
+        // in HOC chains like _s(hoc1(_s(hoc2(_s(actualFunction))))).
+        if (!savedType) {
+          // We're in the innermost call, so this is the actual type.
+          savedType = type;
+          hasCustomHooks = typeof getCustomHooks === 'function';
+        }
+        // Set the signature for all types (even wrappers!) in case
+        // they have no signatures of their own. This is to prevent
+        // problems like https://github.com/facebook/react/issues/20417.
+        if (
+          type != null &&
+          (typeof type === 'function' || typeof type === 'object')
+        ) {
+          setSignature(type, key, forceReset, getCustomHooks);
+        }
+        return type;
+      } else {
+        // We're in the _s() call without arguments, which means
+        // this is the time to collect custom Hook signatures.
+        // Only do this once. This path is hot and runs *inside* every render!
+        if (!didCollectHooks && hasCustomHooks) {
+          didCollectHooks = true;
+          collectCustomHooksForSignature(savedType);
+        }
       }
-      return type;
     };
   } else {
     throw new Error(

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

@@ -524,4 +524,16 @@ describe('ReactFreshBabelPlugin', () => {
         '". If you want to override this check, pass {skipEnvCheck: true} as plugin options.',
     );
   });
+
+  it('does not get tripped by IIFEs', () => {
+    expect(
+      transform(`
+        while (item) {
+          (item => {
+            useFoo();
+          })(item);
+        }
+      `),
+    ).toMatchSnapshot();
+  });
 });

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

@@ -469,6 +469,227 @@ describe('ReactFreshIntegration', () => {
       }
     });
 
+    it('resets state when renaming a state variable inside a HOC with direct call', () => {
+      if (__DEV__) {
+        render(`
+          const {useState} = React;
+          const S = 1;
+
+          function hocWithDirectCall(Wrapped) {
+            return function Generated() {
+              return Wrapped();
+            };
+          }
+
+          export default hocWithDirectCall(() => {
+            const [foo, setFoo] = useState(S);
+            return <h1>A{foo}</h1>;
+          });
+        `);
+        const el = container.firstChild;
+        expect(el.textContent).toBe('A1');
+
+        patch(`
+          const {useState} = React;
+          const S = 2;
+
+          function hocWithDirectCall(Wrapped) {
+            return function Generated() {
+              return Wrapped();
+            };
+          }
+
+          export default hocWithDirectCall(() => {
+            const [foo, setFoo] = useState(S);
+            return <h1>B{foo}</h1>;
+          });
+        `);
+        // Same state variable name, so state is preserved.
+        expect(container.firstChild).toBe(el);
+        expect(el.textContent).toBe('B1');
+
+        patch(`
+          const {useState} = React;
+          const S = 3;
+
+          function hocWithDirectCall(Wrapped) {
+            return function Generated() {
+              return Wrapped();
+            };
+          }
+
+          export default hocWithDirectCall(() => {
+            const [bar, setBar] = useState(S);
+            return <h1>C{bar}</h1>;
+          });
+        `);
+        // Different state variable name, so state is reset.
+        expect(container.firstChild).not.toBe(el);
+        const newEl = container.firstChild;
+        expect(newEl.textContent).toBe('C3');
+      }
+    });
+
+    it('does not crash when changing Hook order inside a HOC with direct call', () => {
+      if (__DEV__) {
+        render(`
+          const {useEffect} = React;
+
+          function hocWithDirectCall(Wrapped) {
+            return function Generated() {
+              return Wrapped();
+            };
+          }
+
+          export default hocWithDirectCall(() => {
+            useEffect(() => {}, []);
+            return <h1>A</h1>;
+          });
+        `);
+        const el = container.firstChild;
+        expect(el.textContent).toBe('A');
+
+        patch(`
+          const {useEffect} = React;
+
+          function hocWithDirectCall(Wrapped) {
+            return function Generated() {
+              return Wrapped();
+            };
+          }
+
+          export default hocWithDirectCall(() => {
+            useEffect(() => {}, []);
+            useEffect(() => {}, []);
+            return <h1>B</h1>;
+          });
+        `);
+        // Hook order changed, so we remount.
+        expect(container.firstChild).not.toBe(el);
+        const newEl = container.firstChild;
+        expect(newEl.textContent).toBe('B');
+      }
+    });
+
+    it('does not crash when changing Hook order inside a memo-ed HOC with direct call', () => {
+      if (__DEV__) {
+        render(`
+          const {useEffect, memo} = React;
+
+          function hocWithDirectCall(Wrapped) {
+            return memo(function Generated() {
+              return Wrapped();
+            });
+          }
+
+          export default hocWithDirectCall(() => {
+            useEffect(() => {}, []);
+            return <h1>A</h1>;
+          });
+        `);
+        const el = container.firstChild;
+        expect(el.textContent).toBe('A');
+
+        patch(`
+          const {useEffect, memo} = React;
+
+          function hocWithDirectCall(Wrapped) {
+            return memo(function Generated() {
+              return Wrapped();
+            });
+          }
+
+          export default hocWithDirectCall(() => {
+            useEffect(() => {}, []);
+            useEffect(() => {}, []);
+            return <h1>B</h1>;
+          });
+        `);
+        // Hook order changed, so we remount.
+        expect(container.firstChild).not.toBe(el);
+        const newEl = container.firstChild;
+        expect(newEl.textContent).toBe('B');
+      }
+    });
+
+    it('does not crash when changing Hook order inside a memo+forwardRef-ed HOC with direct call', () => {
+      if (__DEV__) {
+        render(`
+          const {useEffect, memo, forwardRef} = React;
+
+          function hocWithDirectCall(Wrapped) {
+            return memo(forwardRef(function Generated() {
+              return Wrapped();
+            }));
+          }
+
+          export default hocWithDirectCall(() => {
+            useEffect(() => {}, []);
+            return <h1>A</h1>;
+          });
+        `);
+        const el = container.firstChild;
+        expect(el.textContent).toBe('A');
+
+        patch(`
+          const {useEffect, memo, forwardRef} = React;
+
+          function hocWithDirectCall(Wrapped) {
+            return memo(forwardRef(function Generated() {
+              return Wrapped();
+            }));
+          }
+
+          export default hocWithDirectCall(() => {
+            useEffect(() => {}, []);
+            useEffect(() => {}, []);
+            return <h1>B</h1>;
+          });
+        `);
+        // Hook order changed, so we remount.
+        expect(container.firstChild).not.toBe(el);
+        const newEl = container.firstChild;
+        expect(newEl.textContent).toBe('B');
+      }
+    });
+
+    it('does not crash when changing Hook order inside a HOC returning an object', () => {
+      if (__DEV__) {
+        render(`
+          const {useEffect} = React;
+
+          function hocWithDirectCall(Wrapped) {
+            return {Wrapped: Wrapped};
+          }
+
+          export default hocWithDirectCall(() => {
+            useEffect(() => {}, []);
+            return <h1>A</h1>;
+          }).Wrapped;
+        `);
+        const el = container.firstChild;
+        expect(el.textContent).toBe('A');
+
+        patch(`
+          const {useEffect} = React;
+
+          function hocWithDirectCall(Wrapped) {
+            return {Wrapped: Wrapped};
+          }
+
+          export default hocWithDirectCall(() => {
+            useEffect(() => {}, []);
+            useEffect(() => {}, []);
+            return <h1>B</h1>;
+          }).Wrapped;
+        `);
+        // Hook order changed, so we remount.
+        expect(container.firstChild).not.toBe(el);
+        const newEl = container.firstChild;
+        expect(newEl.textContent).toBe('B');
+      }
+    });
+
     it('resets effects while preserving state', () => {
       if (__DEV__) {
         render(`

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

@@ -37,11 +37,13 @@ const Bar = () => {
 _s4(Bar, "useContext{}");
 
 _c2 = Bar;
-const Baz = memo(_c3 = _s5(() => {
+
+const Baz = _s5(memo(_c3 = _s5(() => {
   _s5();
 
   return useContext(X);
-}, "useContext{}"));
+}, "useContext{}")), "useContext{}");
+
 _c4 = Baz;
 
 const Qux = () => {
@@ -84,6 +86,18 @@ var _c;
 $RefreshReg$(_c, "App");
 `;
 
+exports[`ReactFreshBabelPlugin does not get tripped by IIFEs 1`] = `
+while (item) {
+  var _s = $RefreshSig$();
+
+  _s(item => {
+    _s();
+
+    useFoo();
+  }, "useFoo{}", true)(item);
+}
+`;
+
 exports[`ReactFreshBabelPlugin generates signatures for function declarations calling hooks 1`] = `
 var _s = $RefreshSig$();
 
@@ -108,21 +122,21 @@ exports[`ReactFreshBabelPlugin generates signatures for function expressions cal
 var _s = $RefreshSig$(),
     _s2 = $RefreshSig$();
 
-export const A = React.memo(_c2 = React.forwardRef(_c = _s((props, ref) => {
+export const A = _s(React.memo(_c2 = _s(React.forwardRef(_c = _s((props, ref) => {
   _s();
 
   const [foo, setFoo] = useState(0);
   React.useEffect(() => {});
   return <h1 ref={ref}>{foo}</h1>;
-}, "useState{[foo, setFoo](0)}\\nuseEffect{}")));
+}, "useState{[foo, setFoo](0)}\\nuseEffect{}")), "useState{[foo, setFoo](0)}\\nuseEffect{}")), "useState{[foo, setFoo](0)}\\nuseEffect{}");
 _c3 = A;
-export const B = React.memo(_c5 = React.forwardRef(_c4 = _s2(function (props, ref) {
+export const B = _s2(React.memo(_c5 = _s2(React.forwardRef(_c4 = _s2(function (props, ref) {
   _s2();
 
   const [foo, setFoo] = useState(0);
   React.useEffect(() => {});
   return <h1 ref={ref}>{foo}</h1>;
-}, "useState{[foo, setFoo](0)}\\nuseEffect{}")));
+}, "useState{[foo, setFoo](0)}\\nuseEffect{}")), "useState{[foo, setFoo](0)}\\nuseEffect{}")), "useState{[foo, setFoo](0)}\\nuseEffect{}");
 _c6 = B;
 
 function hoc() {