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';
@@ -292,9 +296,6 @@ function propagateContextChange_eager<T>(
  }
  dependency = dependency.next;
  }
- } else if (fiber.tag === ContextProvider) {
- // Don't scan deeper if this is a matching provider
- nextFiber = fiber.type === workInProgress.type ? null : fiber.child;
  } else if (
  enableSuspenseServerRenderer &&
  fiber.tag === DehydratedFragment
@@ -409,14 +410,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 +506,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 +522,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 +584,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 +706,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;