Specify propagation root for Suspense too

facebook · Jan 18, 2022 · 28d82de · 28d82de
1 parent 92c5c04
commit 28d82de
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.new.js b/packages/react-reconciler/src/ReactFiberBeginWork.new.js
@@ -2754,13 +2754,17 @@ function updateDehydratedSuspenseComponent(
   }
 }
 
-function scheduleContextWorkOnFiber(fiber: Fiber, renderLanes: Lanes) {
+function scheduleSuspenseWorkOnFiber(fiber: Fiber, renderLanes: Lanes) {
   fiber.lanes = mergeLanes(fiber.lanes, renderLanes);
   const alternate = fiber.alternate;
   if (alternate !== null) {
     alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
   }
-  scheduleContextWorkOnParentPath(fiber.return, renderLanes, null);
+  scheduleContextWorkOnParentPath(
+    fiber.return,
+    renderLanes,
+    fiber.return.return,
+  );
 }
 
 function propagateSuspenseContextChange(
@@ -2776,15 +2780,15 @@ function propagateSuspenseContextChange(
     if (node.tag === SuspenseComponent) {
       const state: SuspenseState | null = node.memoizedState;
       if (state !== null) {
-        scheduleContextWorkOnFiber(node, renderLanes);
+        scheduleSuspenseWorkOnFiber(node, renderLanes);
       }
     } else if (node.tag === SuspenseListComponent) {
       // If the tail is hidden there might not be an Suspense boundaries
       // to schedule work on. In this case we have to schedule it on the
       // list itself.
       // We don't have to traverse to the children of the list since
       // the list will propagate the change when it rerenders.
-      scheduleContextWorkOnFiber(node, renderLanes);
+      scheduleSuspenseWorkOnFiber(node, renderLanes);
     } else if (node.child !== null) {
       node.child.return = node;
       node = node.child;

diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.old.js b/packages/react-reconciler/src/ReactFiberBeginWork.old.js
@@ -2754,13 +2754,17 @@ function updateDehydratedSuspenseComponent(
   }
 }
 
-function scheduleContextWorkOnFiber(fiber: Fiber, renderLanes: Lanes) {
+function scheduleSuspenseWorkOnFiber(fiber: Fiber, renderLanes: Lanes) {
   fiber.lanes = mergeLanes(fiber.lanes, renderLanes);
   const alternate = fiber.alternate;
   if (alternate !== null) {
     alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
   }
-  scheduleContextWorkOnParentPath(fiber.return, renderLanes, null);
+  scheduleContextWorkOnParentPath(
+    fiber.return,
+    renderLanes,
+    fiber.return.return,
+  );
 }
 
 function propagateSuspenseContextChange(
@@ -2776,15 +2780,15 @@ function propagateSuspenseContextChange(
     if (node.tag === SuspenseComponent) {
       const state: SuspenseState | null = node.memoizedState;
       if (state !== null) {
-        scheduleContextWorkOnFiber(node, renderLanes);
+        scheduleSuspenseWorkOnFiber(node, renderLanes);
       }
     } else if (node.tag === SuspenseListComponent) {
       // If the tail is hidden there might not be an Suspense boundaries
       // to schedule work on. In this case we have to schedule it on the
       // list itself.
       // We don't have to traverse to the children of the list since
       // the list will propagate the change when it rerenders.
-      scheduleContextWorkOnFiber(node, renderLanes);
+      scheduleSuspenseWorkOnFiber(node, renderLanes);
     } else if (node.child !== null) {
       node.child.return = node;
       node = node.child;