Attach signatures at every nesting level

packages/react-refresh/src/ReactFreshBabelPlugin.js

@@ -333,24 +333,23 @@ export default function(babel, opts = {}) {
     return args;
   }
 
-  // Traverse HOC calls upwards to the rootmost one.
-  function findOuterCallPath(path) {
-    let outerCallPath = null;
+  function findHOCCallPathsAbove(path) {
+    let calls = [];
     while (true) {
       if (!path) {
-        return outerCallPath;
+        return calls;
       }
       if (path.node.type === 'AssignmentExpression') {
         // Ignore registrations.
         path = path.parentPath;
         continue;
       }
       if (path.node.type === 'CallExpression') {
-        outerCallPath = path;
+        calls.push(path);
         path = path.parentPath;
         continue;
       }
-      return outerCallPath; // Stop at other types.
+      return calls; // Stop at other types.
     }
   }
 
@@ -651,17 +650,16 @@ export default function(babel, opts = {}) {
             // Result: let Foo = () => {}; __signature(Foo, ...);
           } else {
             // let Foo = hoc(() => {})
-            const outerCallPath = findOuterCallPath(path.parentPath);
-            if (outerCallPath) {
-              path = outerCallPath;
-            }
-            path.replaceWith(
-              t.callExpression(
-                sigCallID,
-                createArgumentsForSignature(path.node, signature, path.scope),
-              ),
-            );
-            // Result: let Foo = __signature(hoc(() => {}), ...)
+            const paths = [path, ...findHOCCallPathsAbove(path.parentPath)];
+            paths.forEach(path => {
+              path.replaceWith(
+                t.callExpression(
+                  sigCallID,
+                  createArgumentsForSignature(path.node, signature, path.scope),
+                ),
+              );
+            });
+            // Result: let Foo = __signature(hoc(__signature(() => {}, ...)), ...)
           }
         },
       },

packages/react-refresh/src/ReactFreshRuntime.js

@@ -612,57 +612,53 @@ 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;
+      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.
+        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__/__snapshots__/ReactFreshBabelPlugin-test.js.snap

@@ -38,11 +38,11 @@ _s4(Bar, "useContext{}");
 
 _c2 = Bar;
 
-const Baz = _s5(memo(_c3 = () => {
+const Baz = _s5(memo(_c3 = _s5(() => {
   _s5();
 
   return useContext(X);
-}), "useContext{}");
+}, "useContext{}")), "useContext{}");
 
 _c4 = Baz;
 
@@ -110,21 +110,21 @@ exports[`ReactFreshBabelPlugin generates signatures for function expressions cal
 var _s = $RefreshSig$(),
     _s2 = $RefreshSig$();
 
-export const A = _s(React.memo(_c2 = React.forwardRef(_c = (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 = _s2(React.memo(_c5 = React.forwardRef(_c4 = 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() {