Specify propagation root for Suspense too

Author: gaearon

packages/react-reconciler/src/ReactFiberBeginWork.new.js

@@ -2754,13 +2754,15 @@ 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);
+  // Guaranteed to be non-empty since the fiber is not a root.
+  const parentFiber: Fiber = (fiber.return: any);
+  scheduleContextWorkOnParentPath(parentFiber, renderLanes, parentFiber);
 }
 
 function propagateSuspenseContextChange(
@@ -2776,15 +2778,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;

packages/react-reconciler/src/ReactFiberBeginWork.old.js

@@ -2754,13 +2754,15 @@ 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);
+  // Guaranteed to be non-empty since the fiber is not a root.
+  const parentFiber: Fiber = (fiber.return: any);
+  scheduleContextWorkOnParentPath(parentFiber, renderLanes, parentFiber);
 }
 
 function propagateSuspenseContextChange(
@@ -2776,15 +2778,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;

packages/react-reconciler/src/ReactFiberNewContext.new.js

@@ -141,7 +141,7 @@ export function popProvider(
 export function scheduleContextWorkOnParentPath(
   parent: Fiber | null,
   renderLanes: Lanes,
-  propagationRoot: Fiber | null,
+  propagationRoot: Fiber,
 ) {
   // Update the child lanes of all the ancestors, including the alternates.
   let node = parent;
@@ -169,7 +169,7 @@ export function scheduleContextWorkOnParentPath(
       break;
     }
     if (__DEV__) {
-      if (propagationRoot !== null && node === propagationRoot.alternate) {
+      if (node === propagationRoot.alternate) {
         console.error(
           'Did not expect to encounter a propagation root alternate when scheduling context work. ' +
             'This error is likely caused by a bug in React. Please file an issue.',

packages/react-reconciler/src/ReactFiberNewContext.old.js

@@ -141,7 +141,7 @@ export function popProvider(
 export function scheduleContextWorkOnParentPath(
   parent: Fiber | null,
   renderLanes: Lanes,
-  propagationRoot: Fiber | null,
+  propagationRoot: Fiber,
 ) {
   // Update the child lanes of all the ancestors, including the alternates.
   let node = parent;
@@ -169,7 +169,7 @@ export function scheduleContextWorkOnParentPath(
       break;
     }
     if (__DEV__) {
-      if (propagationRoot !== null && node === propagationRoot.alternate) {
+      if (node === propagationRoot.alternate) {
         console.error(
           'Did not expect to encounter a propagation root alternate when scheduling context work. ' +
             'This error is likely caused by a bug in React. Please file an issue.',