From ec38def44f2ff696ec0b73c212265fd2324bcdaf Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Mon, 20 May 2019 14:55:01 +0100
Subject: [PATCH] [Fresh] Don't traverse remounted trees (#15685)

* Don't traverse children when hot reloading needs a remount

If we're gonna remount that tree anyway, there is no use in traversing its children beforehand.

* Add a test verifying hot reload batches updates

Otherwise there is a risk of it being super slow due to cascades.
---
 .../src/__tests__/ReactFresh-test.internal.js | 41 +++++++++++++++++++
 .../src/ReactFiberHotReloading.js             | 19 +++++----
 2 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactFresh-test.internal.js b/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
index 6cf4189854..ec5c733494 100644
--- a/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
+++ b/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
@@ -1549,6 +1549,47 @@ describe('ReactFresh', () => {
     }
   });
 
+  it('batches re-renders during a hot update', () => {
+    if (__DEV__) {
+      let helloRenders = 0;
+
+      render(() => {
+        function Hello({children}) {
+          helloRenders++;
+          return <div>X{children}X</div>;
+        }
+        __register__(Hello, 'Hello');
+
+        function App() {
+          return (
+            <Hello>
+              <Hello>
+                <Hello />
+              </Hello>
+              <Hello>
+                <Hello />
+              </Hello>
+            </Hello>
+          );
+        }
+        return App;
+      });
+      expect(helloRenders).toBe(5);
+      expect(container.textContent).toBe('XXXXXXXXXX');
+      helloRenders = 0;
+
+      patch(() => {
+        function Hello({children}) {
+          helloRenders++;
+          return <div>O{children}O</div>;
+        }
+        __register__(Hello, 'Hello');
+      });
+      expect(helloRenders).toBe(5);
+      expect(container.textContent).toBe('OOOOOOOOOO');
+    }
+  });
+
   it('does not leak state between components', () => {
     if (__DEV__) {
       const AppV1 = render(
diff --git a/packages/react-reconciler/src/ReactFiberHotReloading.js b/packages/react-reconciler/src/ReactFiberHotReloading.js
index 1804ba6fa5..c6a7d2e0db 100644
--- a/packages/react-reconciler/src/ReactFiberHotReloading.js
+++ b/packages/react-reconciler/src/ReactFiberHotReloading.js
@@ -224,29 +224,34 @@ function scheduleFibersWithFamiliesRecursively(
       throw new Error('Expected familiesByType to be set during hot reload.');
     }
 
+    let needsRender = false;
+    let needsRemount = false;
     if (candidateType !== null) {
       const family = familiesByType.get(candidateType);
       if (family !== undefined) {
         if (staleFamilies.has(family)) {
-          fiber._debugNeedsRemount = true;
-          scheduleWork(fiber, Sync);
+          needsRemount = true;
         } else if (updatedFamilies.has(family)) {
-          scheduleWork(fiber, Sync);
+          needsRender = true;
         }
       }
     }
-
     if (failedBoundaries !== null) {
       if (
         failedBoundaries.has(fiber) ||
         (alternate !== null && failedBoundaries.has(alternate))
       ) {
-        fiber._debugNeedsRemount = true;
-        scheduleWork(fiber, Sync);
+        needsRemount = true;
       }
     }
 
-    if (child !== null) {
+    if (needsRemount) {
+      fiber._debugNeedsRemount = true;
+    }
+    if (needsRemount || needsRender) {
+      scheduleWork(fiber, Sync);
+    }
+    if (child !== null && !needsRemount) {
       scheduleFibersWithFamiliesRecursively(
         child,
         updatedFamilies,