[Fresh] Support re-rendering lazy() without losing state (#15686)

* Support re-rendering lazy() without losing state

* Clearer naming

packages/react-dom/src/__tests__/ReactFresh-test.internal.js

@@ -897,7 +897,7 @@ describe('ReactFresh', () => {
 
   it('can preserve state for lazy after resolution', async () => {
     if (__DEV__) {
-      let OuterV1 = render(() => {
+      let AppV1 = render(() => {
         function Hello() {
           const [val, setVal] = React.useState(0);
           return (
@@ -923,6 +923,7 @@ describe('ReactFresh', () => {
             </React.Suspense>
           );
         }
+        __register__(App, 'App');
 
         return App;
       });
@@ -943,7 +944,7 @@ describe('ReactFresh', () => {
       expect(el.textContent).toBe('1');
 
       // Perform a hot update.
-      patch(() => {
+      let AppV2 = patch(() => {
         function Hello() {
           const [val, setVal] = React.useState(0);
           return (
@@ -953,6 +954,25 @@ describe('ReactFresh', () => {
           );
         }
         __register__(Hello, 'Hello');
+
+        const Outer = React.lazy(
+          () =>
+            new Promise(resolve => {
+              setTimeout(() => resolve({default: Hello}), 100);
+            }),
+        );
+        __register__(Outer, 'Outer');
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Outer />
+            </React.Suspense>
+          );
+        }
+        __register__(App, 'App');
+
+        return App;
       });
 
       // Assert the state was preserved but color changed.
@@ -968,8 +988,12 @@ describe('ReactFresh', () => {
       expect(el.textContent).toBe('2');
       expect(el.style.color).toBe('red');
 
-      // Perform a top-down render with a stale type.
-      render(() => OuterV1);
+      // Perform top-down renders with both fresh and stale types.
+      // Neither should change the state or color.
+      // They should always resolve to the latest version.
+      render(() => AppV1);
+      render(() => AppV2);
+      render(() => AppV1);
       expect(container.firstChild).toBe(el);
       expect(el.textContent).toBe('2');
       expect(el.style.color).toBe('red');
@@ -987,7 +1011,17 @@ describe('ReactFresh', () => {
         __register__(Hello, 'Hello');
 
         // Note: no lazy wrapper this time.
-        return Hello;
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Hello />
+            </React.Suspense>
+          );
+        }
+        __register__(App, 'App');
+
+        return App;
       });
 
       expect(container.firstChild).not.toBe(el);
@@ -2053,6 +2087,18 @@ describe('ReactFresh', () => {
     }
   });
 
+  it('can remount on signature change within a lazy simple memo wrapper', () => {
+    if (__DEV__) {
+      testRemountingWithWrapper(Hello =>
+        React.lazy(() => ({
+          then(cb) {
+            cb({default: React.memo(Hello)});
+          },
+        })),
+      );
+    }
+  });
+
   it('can remount on signature change within forwardRef', () => {
     if (__DEV__) {
       testRemountingWithWrapper(Hello => React.forwardRef(Hello));

packages/react-reconciler/src/ReactFiberHotReloading.js

@@ -23,7 +23,11 @@ import {
   MemoComponent,
   SimpleMemoComponent,
 } from 'shared/ReactWorkTags';
-import {REACT_FORWARD_REF_TYPE, REACT_MEMO_TYPE} from 'shared/ReactSymbols';
+import {
+  REACT_FORWARD_REF_TYPE,
+  REACT_MEMO_TYPE,
+  REACT_LAZY_TYPE,
+} from 'shared/ReactSymbols';
 
 type Family = {|
   currentType: any,
@@ -108,39 +112,47 @@ export function isCompatibleFamilyForHotReloading(
 
     const prevType = fiber.elementType;
     const nextType = element.type;
+
     // If we got here, we know types aren't === equal.
     let needsCompareFamilies = false;
+
+    const $$typeofNextType =
+      typeof nextType === 'object' && nextType !== null
+        ? nextType.$$typeof
+        : null;
+
     switch (fiber.tag) {
       case FunctionComponent: {
         if (typeof nextType === 'function') {
           needsCompareFamilies = true;
+        } else if ($$typeofNextType === REACT_LAZY_TYPE) {
+          // We don't know the inner type yet.
+          // We're going to assume that the lazy inner type is stable,
+          // and so it is sufficient to avoid reconciling it away.
+          // We're not going to unwrap or actually use the new lazy type.
+          needsCompareFamilies = true;
         }
         break;
       }
       case ForwardRef: {
-        if (
-          typeof nextType === 'object' &&
-          nextType !== null &&
-          nextType.$$typeof === REACT_FORWARD_REF_TYPE
-        ) {
+        if ($$typeofNextType === REACT_FORWARD_REF_TYPE) {
+          needsCompareFamilies = true;
+        } else if ($$typeofNextType === REACT_LAZY_TYPE) {
           needsCompareFamilies = true;
         }
         break;
       }
       case MemoComponent:
       case SimpleMemoComponent: {
-        if (
-          typeof nextType === 'object' &&
-          nextType !== null &&
-          nextType.$$typeof === REACT_MEMO_TYPE
-        ) {
+        if ($$typeofNextType === REACT_MEMO_TYPE) {
           // TODO: if it was but can no longer be simple,
           // we shouldn't set this.
           needsCompareFamilies = true;
+        } else if ($$typeofNextType === REACT_LAZY_TYPE) {
+          needsCompareFamilies = true;
         }
         break;
       }
-      // TODO: maybe support lazy?
       default:
         return false;
     }