Sign nested memo/forwardRef too

facebook · Mar 26, 2021 · 77a9ef3 · 77a9ef3
1 parent a804fbc
commit 77a9ef3
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 8 deletions.
diff --git a/packages/react-refresh/src/ReactFreshRuntime.js b/packages/react-refresh/src/ReactFreshRuntime.js
@@ -355,15 +355,25 @@ export function setSignature(
   getCustomHooks?: () => Array<Function>,
 ): void {
   if (__DEV__) {
-    if (allSignaturesByType.has(type)) {
-      return;
+    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;
+      }
     }
-    allSignaturesByType.set(type, {
-      forceReset,
-      ownKey: key,
-      fullKey: null,
-      getCustomHooks: getCustomHooks || (() => []),
-    });
   } else {
     throw new Error(
       'Unexpected call to React Refresh in a production environment.',

diff --git a/packages/react-refresh/src/__tests__/ReactFreshIntegration-test.js b/packages/react-refresh/src/__tests__/ReactFreshIntegration-test.js
@@ -571,6 +571,88 @@ describe('ReactFreshIntegration', () => {
       }
     });
 
+    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(`