Store list of contexts on the fiber

Currently, context can only be read by a special type of component,
ContextConsumer. We want to add support to all fibers, including
classes and functional components.

Each fiber may read from one or more contexts. To enable quick, mono-
morphic access of this list, we'll store them on a fiber property.

Author: acdlite

packages/react-reconciler/src/ReactFiber.js

@@ -14,6 +14,7 @@ import type {TypeOfMode} from './ReactTypeOfMode';
 import type {TypeOfSideEffect} from 'shared/ReactTypeOfSideEffect';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {UpdateQueue} from './ReactUpdateQueue';
+import type {ContextReader} from './ReactFiberNewContext';
 
 import invariant from 'shared/invariant';
 import {enableProfilerTimer} from 'shared/ReactFeatureFlags';
@@ -124,6 +125,9 @@ export type Fiber = {|
   // The state used to create the output
   memoizedState: any,
 
+  // A linked-list of contexts that this fiber depends on
+  firstContextReader: ContextReader<mixed> | null,
+
   // Bitfield that describes properties about the fiber and its subtree. E.g.
   // the AsyncMode flag indicates whether the subtree should be async-by-
   // default. When a fiber is created, it inherits the mode of its
@@ -213,6 +217,7 @@ function FiberNode(
   this.memoizedProps = null;
   this.updateQueue = null;
   this.memoizedState = null;
+  this.firstContextReader = null;
 
   this.mode = mode;
 
@@ -331,6 +336,7 @@ export function createWorkInProgress(
   workInProgress.memoizedProps = current.memoizedProps;
   workInProgress.memoizedState = current.memoizedState;
   workInProgress.updateQueue = current.updateQueue;
+  workInProgress.firstContextReader = current.firstContextReader;
 
   // These will be overridden during the parent's reconciliation
   workInProgress.sibling = current.sibling;
@@ -562,6 +568,7 @@ export function assignFiberPropertiesInDEV(
   target.memoizedProps = source.memoizedProps;
   target.updateQueue = source.updateQueue;
   target.memoizedState = source.memoizedState;
+  target.firstContextReader = source.firstContextReader;
   target.mode = source.mode;
   target.effectTag = source.effectTag;
   target.nextEffect = source.nextEffect;

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -69,8 +69,11 @@ import {
 import {pushHostContext, pushHostContainer} from './ReactFiberHostContext';
 import {
   pushProvider,
-  getContextCurrentValue,
-  getContextChangedBits,
+  propagateContextChange,
+  checkForPendingContext,
+  readContext,
+  prepareToReadContext,
+  finishReadingContext,
 } from './ReactFiberNewContext';
 import {
   markActualRenderTimeStarted,
@@ -764,100 +767,6 @@ function updatePortalComponent(current, workInProgress, renderExpirationTime) {
   return workInProgress.child;
 }
 
-function propagateContextChange<V>(
-  workInProgress: Fiber,
-  context: ReactContext<V>,
-  changedBits: number,
-  renderExpirationTime: ExpirationTime,
-): void {
-  let fiber = workInProgress.child;
-  if (fiber !== null) {
-    // Set the return pointer of the child to the work-in-progress fiber.
-    fiber.return = workInProgress;
-  }
-  while (fiber !== null) {
-    let nextFiber;
-    // Visit this fiber.
-    switch (fiber.tag) {
-      case ContextConsumer:
-        // Check if the context matches.
-        const observedBits: number = fiber.stateNode | 0;
-        if (fiber.type === context && (observedBits & changedBits) !== 0) {
-          // Update the expiration time of all the ancestors, including
-          // the alternates.
-          let node = fiber;
-          while (node !== null) {
-            const alternate = node.alternate;
-            if (
-              node.expirationTime === NoWork ||
-              node.expirationTime > renderExpirationTime
-            ) {
-              node.expirationTime = renderExpirationTime;
-              if (
-                alternate !== null &&
-                (alternate.expirationTime === NoWork ||
-                  alternate.expirationTime > renderExpirationTime)
-              ) {
-                alternate.expirationTime = renderExpirationTime;
-              }
-            } else if (
-              alternate !== null &&
-              (alternate.expirationTime === NoWork ||
-                alternate.expirationTime > renderExpirationTime)
-            ) {
-              alternate.expirationTime = renderExpirationTime;
-            } else {
-              // Neither alternate was updated, which means the rest of the
-              // ancestor path already has sufficient priority.
-              break;
-            }
-            node = node.return;
-          }
-          // Don't scan deeper than a matching consumer. When we render the
-          // consumer, we'll continue scanning from that point. This way the
-          // scanning work is time-sliced.
-          nextFiber = null;
-        } else {
-          // Traverse down.
-          nextFiber = fiber.child;
-        }
-        break;
-      case ContextProvider:
-        // Don't scan deeper if this is a matching provider
-        nextFiber = fiber.type === workInProgress.type ? null : fiber.child;
-        break;
-      default:
-        // Traverse down.
-        nextFiber = fiber.child;
-        break;
-    }
-    if (nextFiber !== null) {
-      // Set the return pointer of the child to the work-in-progress fiber.
-      nextFiber.return = fiber;
-    } else {
-      // No child. Traverse to next sibling.
-      nextFiber = fiber;
-      while (nextFiber !== null) {
-        if (nextFiber === workInProgress) {
-          // We're back to the root of this subtree. Exit.
-          nextFiber = null;
-          break;
-        }
-        let sibling = nextFiber.sibling;
-        if (sibling !== null) {
-          // Set the return pointer of the sibling to the work-in-progress fiber.
-          sibling.return = nextFiber.return;
-          nextFiber = sibling;
-          break;
-        }
-        // No more siblings. Traverse up.
-        nextFiber = nextFiber.return;
-      }
-    }
-    fiber = nextFiber;
-  }
-}
-
 function updateContextProvider(current, workInProgress, renderExpirationTime) {
   const providerType: ReactProviderType<any> = workInProgress.type;
   const context: ReactContext<any> = providerType._context;
@@ -970,42 +879,16 @@ function updateContextConsumer(current, workInProgress, renderExpirationTime) {
   const newProps = workInProgress.pendingProps;
   const oldProps = workInProgress.memoizedProps;
 
-  const newValue = getContextCurrentValue(context);
-  const changedBits = getContextChangedBits(context);
-
-  if (hasLegacyContextChanged()) {
-    // Normally we can bail out on props equality but if context has changed
-    // we don't do the bailout and we have to reuse existing props instead.
-  } else if (changedBits === 0 && oldProps === newProps) {
-    return bailoutOnAlreadyFinishedWork(current, workInProgress);
-  }
-  workInProgress.memoizedProps = newProps;
-
-  let observedBits = newProps.unstable_observedBits;
-  if (observedBits === undefined || observedBits === null) {
-    // Subscribe to all changes by default
-    observedBits = MAX_SIGNED_31_BIT_INT;
+  if (!checkForPendingContext(workInProgress, renderExpirationTime)) {
+    if (hasLegacyContextChanged()) {
+      // Normally we can bail out on props equality but if context has changed
+      // we don't do the bailout and we have to reuse existing props instead.
+    } else if (oldProps === newProps) {
+      return bailoutOnAlreadyFinishedWork(current, workInProgress);
+    }
   }
-  // Store the observedBits on the fiber's stateNode for quick access.
-  workInProgress.stateNode = observedBits;
 
-  if ((changedBits & observedBits) !== 0) {
-    // Context change propagation stops at matching consumers, for time-
-    // slicing. Continue the propagation here.
-    propagateContextChange(
-      workInProgress,
-      context,
-      changedBits,
-      renderExpirationTime,
-    );
-  } else if (oldProps === newProps) {
-    // Skip over a memoized parent with a bitmask bailout even
-    // if we began working on it because of a deeper matching child.
-    return bailoutOnAlreadyFinishedWork(current, workInProgress);
-  }
-  // There is no bailout on `children` equality because we expect people
-  // to often pass a bound method as a child, but it may reference
-  // `this.state` or `this.props` (and thus needs to re-render on `setState`).
+  workInProgress.memoizedProps = newProps;
 
   const render = newProps.children;
 
@@ -1019,6 +902,8 @@ function updateContextConsumer(current, workInProgress, renderExpirationTime) {
     );
   }
 
+  prepareToReadContext();
+  const newValue = readContext(context, newProps.unstable_observedBits);
   let newChildren;
   if (__DEV__) {
     ReactCurrentOwner.current = workInProgress;
@@ -1028,6 +913,7 @@ function updateContextConsumer(current, workInProgress, renderExpirationTime) {
   } else {
     newChildren = render(newValue);
   }
+  workInProgress.firstContextReader = finishReadingContext();
 
   // React DevTools reads this flag.
   workInProgress.effectTag |= PerformedWork;