[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();
+  });
 });