Mark trees that need propagation in readContext

Instead of storing matched context consumers in a Set, we can mark
when a consumer receives an update inside `readContext`.

I hesistated to put anything in this function because it's such a hot
path, but so are bail outs. Fortunately, we only need to set this flag
once, the first time a context is read. So I think it's a reasonable
trade off.

In exchange, propagation is faster because we no longer need to
accumulate a Set of matched consumers, and fiber bailouts are faster
because we don't need to consult that Set. And the code is simpler.

Author: acdlite

packages/react-reconciler/src/ReactFiberFlags.js

@@ -12,50 +12,51 @@ import {enableCreateEventHandleAPI} from 'shared/ReactFeatureFlags';
 export type Flags = number;
 
 // Don't change these two values. They're used by React Dev Tools.
-export const NoFlags = /*                      */ 0b000000000000000000000;
-export const PerformedWork = /*                */ 0b000000000000000000001;
+export const NoFlags = /*                      */ 0b0000000000000000000000;
+export const PerformedWork = /*                */ 0b0000000000000000000001;
 
 // You can change the rest (and add more).
-export const Placement = /*                    */ 0b000000000000000000010;
-export const Update = /*                       */ 0b000000000000000000100;
+export const Placement = /*                    */ 0b0000000000000000000010;
+export const Update = /*                       */ 0b0000000000000000000100;
 export const PlacementAndUpdate = /*           */ Placement | Update;
-export const Deletion = /*                     */ 0b000000000000000001000;
-export const ChildDeletion = /*                */ 0b000000000000000010000;
-export const ContentReset = /*                 */ 0b000000000000000100000;
-export const Callback = /*                     */ 0b000000000000001000000;
-export const DidCapture = /*                   */ 0b000000000000010000000;
-export const Ref = /*                          */ 0b000000000000100000000;
-export const Snapshot = /*                     */ 0b000000000001000000000;
-export const Passive = /*                      */ 0b000000000010000000000;
-export const Hydrating = /*                    */ 0b000000000100000000000;
+export const Deletion = /*                     */ 0b0000000000000000001000;
+export const ChildDeletion = /*                */ 0b0000000000000000010000;
+export const ContentReset = /*                 */ 0b0000000000000000100000;
+export const Callback = /*                     */ 0b0000000000000001000000;
+export const DidCapture = /*                   */ 0b0000000000000010000000;
+export const Ref = /*                          */ 0b0000000000000100000000;
+export const Snapshot = /*                     */ 0b0000000000001000000000;
+export const Passive = /*                      */ 0b0000000000010000000000;
+export const Hydrating = /*                    */ 0b0000000000100000000000;
 export const HydratingAndUpdate = /*           */ Hydrating | Update;
-export const Visibility = /*                   */ 0b000000001000000000000;
+export const Visibility = /*                   */ 0b0000000001000000000000;
 
 export const LifecycleEffectMask = Passive | Update | Callback | Ref | Snapshot;
 
 // Union of all commit flags (flags with the lifetime of a particular commit)
-export const HostEffectMask = /*               */ 0b000000001111111111111;
+export const HostEffectMask = /*               */ 0b0000000001111111111111;
 
 // These are not really side effects, but we still reuse this field.
-export const Incomplete = /*                   */ 0b000000010000000000000;
-export const ShouldCapture = /*                */ 0b000000100000000000000;
+export const Incomplete = /*                   */ 0b0000000010000000000000;
+export const ShouldCapture = /*                */ 0b0000000100000000000000;
 // TODO (effects) Remove this bit once the new reconciler is synced to the old.
-export const PassiveUnmountPendingDev = /*     */ 0b000001000000000000000;
-export const ForceUpdateForLegacySuspense = /* */ 0b000010000000000000000;
-export const DidPropagateContext = /*          */ 0b000100000000000000000;
+export const PassiveUnmountPendingDev = /*     */ 0b0000001000000000000000;
+export const ForceUpdateForLegacySuspense = /* */ 0b0000010000000000000000;
+export const DidPropagateContext = /*          */ 0b0000100000000000000000;
+export const NeedsPropagation = /*             */ 0b0001000000000000000000;
 
 // Static tags describe aspects of a fiber that are not specific to a render,
 // e.g. a fiber uses a passive effect (even if there are no updates on this particular render).
 // This enables us to defer more work in the unmount case,
 // since we can defer traversing the tree during layout to look for Passive effects,
 // and instead rely on the static flag as a signal that there may be cleanup work.
-export const PassiveStatic = /*                */ 0b001000000000000000000;
+export const PassiveStatic = /*                */ 0b0010000000000000000000;
 
 // These flags allow us to traverse to fibers that have effects on mount
 // without traversing the entire tree after every commit for
 // double invoking
-export const MountLayoutDev = /*               */ 0b010000000000000000000;
-export const MountPassiveDev = /*              */ 0b100000000000000000000;
+export const MountLayoutDev = /*               */ 0b0100000000000000000000;
+export const MountPassiveDev = /*              */ 0b1000000000000000000000;
 
 // Groups of flags that are used in the commit phase to skip over trees that
 // don't contain effects, by checking subtreeFlags.

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

@@ -33,7 +33,11 @@ import {
   mergeLanes,
   pickArbitraryLane,
 } from './ReactFiberLane.new';
-import {NoFlags, DidPropagateContext} from './ReactFiberFlags';
+import {
+  NoFlags,
+  DidPropagateContext,
+  NeedsPropagation,
+} from './ReactFiberFlags';
 
 import invariant from 'shared/invariant';
 import is from 'shared/objectIs';
@@ -409,14 +413,6 @@ function propagateContextChanges<T>(
               // visit them during render. We should continue propagating the
               // siblings, though
               nextFiber = null;
-
-              // Keep track of subtrees whose propagation we deferred
-              if (deferredPropagation === null) {
-                deferredPropagation = new Set([consumer]);
-              } else {
-                deferredPropagation.add(consumer);
-              }
-              nextFiber = null;
             }
 
             // Since we already found a match, we can stop traversing the
@@ -513,29 +509,14 @@ export function propagateParentContextChangesToDeferredTree(
   );
 }
 
-// Used by lazy context propagation algorithm. When we find a context dependency
-// match, we don't propagate the changes any further into that fiber's subtree.
-// We add the matched fibers to this set. Later, if something inside that
-// subtree bails out of rendering, the presence of a parent fiber in this Set
-// tells us that we need to continue propagating.
-//
-// This is a set of _current_ fibers, not work-in-progress fibers. That's why
-// it's a set instead of a flag on the fiber.
-let deferredPropagation: Set<Fiber> | null = null;
-
-export function resetDeferredContextPropagation() {
-  // This is called by prepareFreshStack
-  deferredPropagation = null;
-}
-
 function propagateParentContextChanges(
   current: Fiber,
   workInProgress: Fiber,
   renderLanes: Lanes,
   forcePropagateEntireTree: boolean,
 ) {
   if (!enableLazyContextPropagation) {
-    return false;
+    return;
   }
 
   // Collect all the parent providers that changed. Since this is usually small
@@ -544,58 +525,36 @@ function propagateParentContextChanges(
   let parent = workInProgress;
   let isInsidePropagationBailout = false;
   while (parent !== null) {
-    const currentParent = parent.alternate;
-    invariant(
-      currentParent !== null,
-      'Should have a current fiber. This is a bug in React.',
-    );
-
     if (!isInsidePropagationBailout) {
-      if (deferredPropagation === null) {
-        if ((parent.flags & DidPropagateContext) !== NoFlags) {
-          break;
-        }
-      } else {
-        if (currentParent !== null && deferredPropagation.has(currentParent)) {
-          // We're inside a subtree that previously bailed out of propagation.
-          // We must disregard the the DidPropagateContext flag as we continue
-          // searching for parent providers.
-          isInsidePropagationBailout = true;
-          // We know that none of the providers in between the propagation
-          // bailout and the nearest render bailout above that could have
-          // changed. So we can skip those.
-          do {
-            parent = parent.return;
-            invariant(
-              parent !== null,
-              'Expected to find a bailed out fiber. This is a bug in React.',
-            );
-          } while ((parent.flags & DidPropagateContext) === NoFlags);
-        } else if ((parent.flags & DidPropagateContext) !== NoFlags) {
-          break;
-        }
+      if ((parent.flags & NeedsPropagation) !== NoFlags) {
+        isInsidePropagationBailout = true;
+      } else if ((parent.flags & DidPropagateContext) !== NoFlags) {
+        break;
       }
     }
 
     if (parent.tag === ContextProvider) {
-      if (currentParent !== null) {
-        const oldProps = currentParent.memoizedProps;
-        if (oldProps !== null) {
-          const providerType: ReactProviderType<any> = parent.type;
-          const context: ReactContext<any> = providerType._context;
-
-          const newProps = parent.pendingProps;
-          const newValue = newProps.value;
-
-          const oldValue = oldProps.value;
-
-          const changedBits = calculateChangedBits(context, newValue, oldValue);
-          if (changedBits !== 0) {
-            if (contexts !== null) {
-              contexts.push(context, changedBits);
-            } else {
-              contexts = [context, changedBits];
-            }
+      const currentParent = parent.alternate;
+      invariant(
+        currentParent !== null,
+        'Should have a current fiber. This is a bug in React.',
+      );
+      const oldProps = currentParent.memoizedProps;
+      if (oldProps !== null) {
+        const providerType: ReactProviderType<any> = parent.type;
+        const context: ReactContext<any> = providerType._context;
+
+        const newProps = parent.pendingProps;
+        const newValue = newProps.value;
+
+        const oldValue = oldProps.value;
+
+        const changedBits = calculateChangedBits(context, newValue, oldValue);
+        if (changedBits !== 0) {
+          if (contexts !== null) {
+            contexts.push(context, changedBits);
+          } else {
+            contexts = [context, changedBits];
           }
         }
       }
@@ -628,10 +587,10 @@ function propagateParentContextChanges(
   //
   // Unfortunately, though, we need to ignore this flag when we're inside a
   // tree whose context propagation was deferred — that's what the
-  // `deferredPropagation` set is for.
+  // `NeedsPropagation` flag is for.
   //
-  // If we could instead bail out before entering the siblings' beging phase,
-  // then we could remove both `DidPropagateContext` and `deferredPropagation`.
+  // If we could instead bail out before entering the siblings' begin phase,
+  // then we could remove both `DidPropagateContext` and `NeedsPropagation`.
   // Consider this as part of the next refactor to the fiber tree structure.
   workInProgress.flags |= DidPropagateContext;
 }
@@ -750,6 +709,9 @@ export function readContext<T>(
         // TODO: This is an old field. Delete it.
         responders: null,
       };
+      if (enableLazyContextPropagation) {
+        currentlyRenderingFiber.flags |= NeedsPropagation;
+      }
     } else {
       // Append a new context item.
       lastContextDependency = lastContextDependency.next = contextItem;