Move context comparison to consumer

In the lazy context implementation, not all context changes are
propagated from the provider, so we can't rely on the propagation alone
to mark the consumer as dirty. The consumer needs to compare to the
previous value, like we do for state and context.

I added a `memoizedValue` field to the context dependency type. Then in
the consumer, we iterate over the current dependencies to see if
something changed. We only do this iteration after props and state has
already bailed out, so it's a relatively uncommon path, except at the
root of a changed subtree. Alternatively, we could move these
comparisons into `readContext`, but that's a much hotter path, so I
think this is an appropriate trade off.

packages/react-dom/src/__tests__/ReactLegacyContextDisabled-test.internal.js

@@ -226,7 +226,15 @@ describe('ReactLegacyContextDisabled', () => {
       container,
     );
     expect(container.textContent).toBe('bbb');
-    expect(lifecycleContextLog).toEqual(['b', 'b']); // sCU skipped due to changed context value.
+    if (gate(flags => flags.enableLazyContextPropagation)) {
+      // In the lazy propagation implementation, we don't check if context
+      // changed until after shouldComponentUpdate is run.
+      expect(lifecycleContextLog).toEqual(['b', 'b', 'b']);
+    } else {
+      // In the eager implementation, a dirty flag was set when the parent
+      // changed, so we skipped sCU.
+      expect(lifecycleContextLog).toEqual(['b', 'b']);
+    }
     ReactDOM.unmountComponentAtNode(container);
   });
 });

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

@@ -3074,6 +3074,10 @@ export function markWorkInProgressReceivedUpdate() {
   didReceiveUpdate = true;
 }
 
+export function checkIfWorkInProgressReceivedUpdate() {
+  return didReceiveUpdate;
+}
+
 function bailoutOnAlreadyFinishedWork(
   current: Fiber | null,
   workInProgress: Fiber,

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

@@ -3074,6 +3074,10 @@ export function markWorkInProgressReceivedUpdate() {
   didReceiveUpdate = true;
 }
 
+export function checkIfWorkInProgressReceivedUpdate() {
+  return didReceiveUpdate;
+}
+
 function bailoutOnAlreadyFinishedWork(
   current: Fiber | null,
   workInProgress: Fiber,

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

@@ -20,6 +20,7 @@ import {
   enableSchedulingProfiler,
   warnAboutDeprecatedLifecycles,
   enableStrictEffects,
+  enableLazyContextPropagation,
 } from 'shared/ReactFeatureFlags';
 import ReactStrictModeWarnings from './ReactStrictModeWarnings.new';
 import {isMounted} from './ReactFiberTreeReflection';
@@ -58,7 +59,7 @@ import {
   hasContextChanged,
   emptyContextObject,
 } from './ReactFiberContext.new';
-import {readContext} from './ReactFiberNewContext.new';
+import {readContext, checkIfContextChanged} from './ReactFiberNewContext.new';
 import {
   requestEventTime,
   requestUpdateLane,
@@ -1150,7 +1151,13 @@ function updateClassInstance(
     unresolvedOldProps === unresolvedNewProps &&
     oldState === newState &&
     !hasContextChanged() &&
-    !checkHasForceUpdateAfterProcessing()
+    !checkHasForceUpdateAfterProcessing() &&
+    !(
+      enableLazyContextPropagation &&
+      current !== null &&
+      current.dependencies !== null &&
+      checkIfContextChanged(current.dependencies)
+    )
   ) {
     // If an update was already in progress, we should schedule an Update
     // effect even though we're bailing out, so that cWU/cDU are called.
@@ -1193,7 +1200,15 @@ function updateClassInstance(
       oldState,
       newState,
       nextContext,
-    );
+    ) ||
+    // TODO: In some cases, we'll end up checking if context has changed twice,
+    // both before and after `shouldComponentUpdate` has been called. Not ideal,
+    // but I'm loath to refactor this function. This only happens for memoized
+    // components so it's not that common.
+    (enableLazyContextPropagation &&
+      current !== null &&
+      current.dependencies !== null &&
+      checkIfContextChanged(current.dependencies));
 
   if (shouldUpdate) {
     // In order to support react-lifecycles-compat polyfilled components,

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

@@ -20,6 +20,7 @@ import {
   enableSchedulingProfiler,
   warnAboutDeprecatedLifecycles,
   enableStrictEffects,
+  enableLazyContextPropagation,
 } from 'shared/ReactFeatureFlags';
 import ReactStrictModeWarnings from './ReactStrictModeWarnings.old';
 import {isMounted} from './ReactFiberTreeReflection';
@@ -58,7 +59,7 @@ import {
   hasContextChanged,
   emptyContextObject,
 } from './ReactFiberContext.old';
-import {readContext} from './ReactFiberNewContext.old';
+import {readContext, checkIfContextChanged} from './ReactFiberNewContext.old';
 import {
   requestEventTime,
   requestUpdateLane,
@@ -1150,7 +1151,13 @@ function updateClassInstance(
     unresolvedOldProps === unresolvedNewProps &&
     oldState === newState &&
     !hasContextChanged() &&
-    !checkHasForceUpdateAfterProcessing()
+    !checkHasForceUpdateAfterProcessing() &&
+    !(
+      enableLazyContextPropagation &&
+      current !== null &&
+      current.dependencies !== null &&
+      checkIfContextChanged(current.dependencies)
+    )
   ) {
     // If an update was already in progress, we should schedule an Update
     // effect even though we're bailing out, so that cWU/cDU are called.
@@ -1193,7 +1200,15 @@ function updateClassInstance(
       oldState,
       newState,
       nextContext,
-    );
+    ) ||
+    // TODO: In some cases, we'll end up checking if context has changed twice,
+    // both before and after `shouldComponentUpdate` has been called. Not ideal,
+    // but I'm loath to refactor this function. This only happens for memoized
+    // components so it's not that common.
+    (enableLazyContextPropagation &&
+      current !== null &&
+      current.dependencies !== null &&
+      checkIfContextChanged(current.dependencies));
 
   if (shouldUpdate) {
     // In order to support react-lifecycles-compat polyfilled components,

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

@@ -30,6 +30,7 @@ import {
   decoupleUpdatePriorityFromScheduler,
   enableUseRefAccessWarning,
   enableStrictEffects,
+  enableLazyContextPropagation,
 } from 'shared/ReactFeatureFlags';
 
 import {
@@ -54,7 +55,7 @@ import {
   higherLanePriority,
   DefaultLanePriority,
 } from './ReactFiberLane.new';
-import {readContext} from './ReactFiberNewContext.new';
+import {readContext, checkIfContextChanged} from './ReactFiberNewContext.new';
 import {HostRoot, CacheComponent} from './ReactWorkTags';
 import {
   Update as UpdateEffect,
@@ -83,7 +84,10 @@ import {
 import invariant from 'shared/invariant';
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
 import is from 'shared/objectIs';
-import {markWorkInProgressReceivedUpdate} from './ReactFiberBeginWork.new';
+import {
+  markWorkInProgressReceivedUpdate,
+  checkIfWorkInProgressReceivedUpdate,
+} from './ReactFiberBeginWork.new';
 import {
   UserBlockingPriority,
   NormalPriority,
@@ -496,6 +500,27 @@ export function renderWithHooks<Props, SecondArg>(
       'early return statement.',
   );
 
+  if (enableLazyContextPropagation) {
+    if (current !== null) {
+      if (!checkIfWorkInProgressReceivedUpdate()) {
+        // If there were no changes to props or state, we need to check if there
+        // was a context change. We didn't already do this because there's no
+        // 1:1 correspondence between dependencies and hooks. Although, because
+        // there almost always is in the common case (`readContext` is an
+        // internal API), we could compare in there. OTOH, we only hit this case
+        // if everything else bails out, so on the whole it might be better to
+        // keep the comparison out of the common path.
+        const currentDependencies = current.dependencies;
+        if (
+          currentDependencies !== null &&
+          checkIfContextChanged(currentDependencies)
+        ) {
+          markWorkInProgressReceivedUpdate();
+        }
+      }
+    }
+  }
+
   return children;
 }
 

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

@@ -30,6 +30,7 @@ import {
   decoupleUpdatePriorityFromScheduler,
   enableUseRefAccessWarning,
   enableStrictEffects,
+  enableLazyContextPropagation,
 } from 'shared/ReactFeatureFlags';
 
 import {
@@ -54,7 +55,7 @@ import {
   higherLanePriority,
   DefaultLanePriority,
 } from './ReactFiberLane.old';
-import {readContext} from './ReactFiberNewContext.old';
+import {readContext, checkIfContextChanged} from './ReactFiberNewContext.old';
 import {HostRoot, CacheComponent} from './ReactWorkTags';
 import {
   Update as UpdateEffect,
@@ -83,7 +84,10 @@ import {
 import invariant from 'shared/invariant';
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
 import is from 'shared/objectIs';
-import {markWorkInProgressReceivedUpdate} from './ReactFiberBeginWork.old';
+import {
+  markWorkInProgressReceivedUpdate,
+  checkIfWorkInProgressReceivedUpdate,
+} from './ReactFiberBeginWork.old';
 import {
   UserBlockingPriority,
   NormalPriority,
@@ -496,6 +500,27 @@ export function renderWithHooks<Props, SecondArg>(
       'early return statement.',
   );
 
+  if (enableLazyContextPropagation) {
+    if (current !== null) {
+      if (!checkIfWorkInProgressReceivedUpdate()) {
+        // If there were no changes to props or state, we need to check if there
+        // was a context change. We didn't already do this because there's no
+        // 1:1 correspondence between dependencies and hooks. Although, because
+        // there almost always is in the common case (`readContext` is an
+        // internal API), we could compare in there. OTOH, we only hit this case
+        // if everything else bails out, so on the whole it might be better to
+        // keep the comparison out of the common path.
+        const currentDependencies = current.dependencies;
+        if (
+          currentDependencies !== null &&
+          checkIfContextChanged(currentDependencies)
+        ) {
+          markWorkInProgressReceivedUpdate();
+        }
+      }
+    }
+  }
+
   return children;
 }
 

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

@@ -8,7 +8,11 @@
  */
 
 import type {ReactContext} from 'shared/ReactTypes';
-import type {Fiber, ContextDependency} from './ReactInternalTypes';
+import type {
+  Fiber,
+  ContextDependency,
+  Dependencies,
+} from './ReactInternalTypes';
 import type {StackCursor} from './ReactFiberStack.new';
 import type {Lanes} from './ReactFiberLane.new';
 import type {SharedQueue} from './ReactUpdateQueue.new';
@@ -34,7 +38,10 @@ import invariant from 'shared/invariant';
 import is from 'shared/objectIs';
 import {createUpdate, ForceUpdate} from './ReactUpdateQueue.new';
 import {markWorkInProgressReceivedUpdate} from './ReactFiberBeginWork.new';
-import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
+import {
+  enableSuspenseServerRenderer,
+  enableLazyContextPropagation,
+} from 'shared/ReactFeatureFlags';
 
 const valueCursor: StackCursor<mixed> = createCursor(null);
 
@@ -210,33 +217,46 @@ export function propagateContextChange<T>(
         ) {
           // Match! Schedule an update on this fiber.
 
-          if (fiber.tag === ClassComponent) {
-            // Schedule a force update on the work-in-progress.
-            const lane = pickArbitraryLane(renderLanes);
-            const update = createUpdate(NoTimestamp, lane);
-            update.tag = ForceUpdate;
-            // TODO: Because we don't have a work-in-progress, this will add the
-            // update to the current fiber, too, which means it will persist even if
-            // this render is thrown away. Since it's a race condition, not sure it's
-            // worth fixing.
+          if (enableLazyContextPropagation) {
+            // In the lazy implemenation, don't mark a dirty flag on the
+            // dependency itself. Not all changes are propagated, so we can't
+            // rely on the propagation function alone to determine whether
+            // something has changed; the consumer will check. In the future,
+            // we could add back a dirty flag as an optimization to avoid
+            // double checking, but until we have selectors it's not really
+            // worth the trouble.
+          } else {
+            if (fiber.tag === ClassComponent) {
+              // Schedule a force update on the work-in-progress.
+              const lane = pickArbitraryLane(renderLanes);
+              const update = createUpdate(NoTimestamp, lane);
+              update.tag = ForceUpdate;
+              // TODO: Because we don't have a work-in-progress, this will add the
+              // update to the current fiber, too, which means it will persist even if
+              // this render is thrown away. Since it's a race condition, not sure it's
+              // worth fixing.
 
-            // Inlined `enqueueUpdate` to remove interleaved update check
-            const updateQueue = fiber.updateQueue;
-            if (updateQueue === null) {
-              // Only occurs if the fiber has been unmounted.
-            } else {
-              const sharedQueue: SharedQueue<any> = (updateQueue: any).shared;
-              const pending = sharedQueue.pending;
-              if (pending === null) {
-                // This is the first update. Create a circular list.
-                update.next = update;
+              // Inlined `enqueueUpdate` to remove interleaved update check
+              const updateQueue = fiber.updateQueue;
+              if (updateQueue === null) {
+                // Only occurs if the fiber has been unmounted.
               } else {
-                update.next = pending.next;
-                pending.next = update;
+                const sharedQueue: SharedQueue<any> = (updateQueue: any).shared;
+                const pending = sharedQueue.pending;
+                if (pending === null) {
+                  // This is the first update. Create a circular list.
+                  update.next = update;
+                } else {
+                  update.next = pending.next;
+                  pending.next = update;
+                }
+                sharedQueue.pending = update;
               }
-              sharedQueue.pending = update;
             }
+            // Mark the updated lanes on the list, too.
+            list.lanes = mergeLanes(list.lanes, renderLanes);
           }
+
           fiber.lanes = mergeLanes(fiber.lanes, renderLanes);
           const alternate = fiber.alternate;
           if (alternate !== null) {
@@ -244,9 +264,6 @@ export function propagateContextChange<T>(
           }
           scheduleWorkOnParentPath(fiber.return, renderLanes);
 
-          // Mark the updated lanes on the list, too.
-          list.lanes = mergeLanes(list.lanes, renderLanes);
-
           // Since we already found a match, we can stop traversing the
           // dependency list.
           break;
@@ -311,6 +328,30 @@ export function propagateContextChange<T>(
   }
 }
 
+export function checkIfContextChanged(currentDependencies: Dependencies) {
+  if (!enableLazyContextPropagation) {
+    return false;
+  }
+  // Iterate over the current dependencies to see if something changed. This
+  // only gets called if props and state has already bailed out, so it's a
+  // relatively uncommon path, except at the root of a changed subtree.
+  // Alternatively, we could move these comparisons into `readContext`, but
+  // that's a much hotter path, so I think this is an appropriate trade off.
+  let dependency = currentDependencies.firstContext;
+  while (dependency !== null) {
+    const context = dependency.context;
+    const newValue = isPrimaryRenderer
+      ? context._currentValue
+      : context._currentValue2;
+    const oldValue = dependency.memoizedValue;
+    if (!is(newValue, oldValue)) {
+      return true;
+    }
+    dependency = dependency.next;
+  }
+  return false;
+}
+
 export function prepareToReadContext(
   workInProgress: Fiber,
   renderLanes: Lanes,
@@ -321,14 +362,19 @@ export function prepareToReadContext(
 
   const dependencies = workInProgress.dependencies;
   if (dependencies !== null) {
-    const firstContext = dependencies.firstContext;
-    if (firstContext !== null) {
-      if (includesSomeLane(dependencies.lanes, renderLanes)) {
-        // Context list has a pending update. Mark that this fiber performed work.
-        markWorkInProgressReceivedUpdate();
-      }
+    if (enableLazyContextPropagation) {
       // Reset the work-in-progress list
       dependencies.firstContext = null;
+    } else {
+      const firstContext = dependencies.firstContext;
+      if (firstContext !== null) {
+        if (includesSomeLane(dependencies.lanes, renderLanes)) {
+          // Context list has a pending update. Mark that this fiber performed work.
+          markWorkInProgressReceivedUpdate();
+        }
+        // Reset the work-in-progress list
+        dependencies.firstContext = null;
+      }
     }
   }
 }
@@ -350,6 +396,10 @@ export function readContext<T>(
     }
   }
 
+  const value = isPrimaryRenderer
+    ? context._currentValue
+    : context._currentValue2;
+
   if (lastContextWithAllBitsObserved === context) {
     // Nothing to do. We already observe everything in this context.
   } else if (observedBits === false || observedBits === 0) {
@@ -370,6 +420,7 @@ export function readContext<T>(
     const contextItem = {
       context: ((context: any): ReactContext<mixed>),
       observedBits: resolvedObservedBits,
+      memoizedValue: value,
       next: null,
     };
 
@@ -387,6 +438,8 @@ export function readContext<T>(
       currentlyRenderingFiber.dependencies = {
         lanes: NoLanes,
         firstContext: contextItem,
+
+        // TODO: This is an old field. Delete it.
         responders: null,
       };
     } else {
@@ -394,5 +447,5 @@ export function readContext<T>(
       lastContextDependency = lastContextDependency.next = contextItem;
     }
   }
-  return isPrimaryRenderer ? context._currentValue : context._currentValue2;
+  return value;
 }

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

@@ -8,7 +8,11 @@
  */
 
 import type {ReactContext} from 'shared/ReactTypes';
-import type {Fiber, ContextDependency} from './ReactInternalTypes';
+import type {
+  Fiber,
+  ContextDependency,
+  Dependencies,
+} from './ReactInternalTypes';
 import type {StackCursor} from './ReactFiberStack.old';
 import type {Lanes} from './ReactFiberLane.old';
 import type {SharedQueue} from './ReactUpdateQueue.old';
@@ -34,7 +38,10 @@ import invariant from 'shared/invariant';
 import is from 'shared/objectIs';
 import {createUpdate, ForceUpdate} from './ReactUpdateQueue.old';
 import {markWorkInProgressReceivedUpdate} from './ReactFiberBeginWork.old';
-import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
+import {
+  enableSuspenseServerRenderer,
+  enableLazyContextPropagation,
+} from 'shared/ReactFeatureFlags';
 
 const valueCursor: StackCursor<mixed> = createCursor(null);
 
@@ -210,33 +217,46 @@ export function propagateContextChange<T>(
         ) {
           // Match! Schedule an update on this fiber.
 
-          if (fiber.tag === ClassComponent) {
-            // Schedule a force update on the work-in-progress.
-            const lane = pickArbitraryLane(renderLanes);
-            const update = createUpdate(NoTimestamp, lane);
-            update.tag = ForceUpdate;
-            // TODO: Because we don't have a work-in-progress, this will add the
-            // update to the current fiber, too, which means it will persist even if
-            // this render is thrown away. Since it's a race condition, not sure it's
-            // worth fixing.
+          if (enableLazyContextPropagation) {
+            // In the lazy implemenation, don't mark a dirty flag on the
+            // dependency itself. Not all changes are propagated, so we can't
+            // rely on the propagation function alone to determine whether
+            // something has changed; the consumer will check. In the future,
+            // we could add back a dirty flag as an optimization to avoid
+            // double checking, but until we have selectors it's not really
+            // worth the trouble.
+          } else {
+            if (fiber.tag === ClassComponent) {
+              // Schedule a force update on the work-in-progress.
+              const lane = pickArbitraryLane(renderLanes);
+              const update = createUpdate(NoTimestamp, lane);
+              update.tag = ForceUpdate;
+              // TODO: Because we don't have a work-in-progress, this will add the
+              // update to the current fiber, too, which means it will persist even if
+              // this render is thrown away. Since it's a race condition, not sure it's
+              // worth fixing.
 
-            // Inlined `enqueueUpdate` to remove interleaved update check
-            const updateQueue = fiber.updateQueue;
-            if (updateQueue === null) {
-              // Only occurs if the fiber has been unmounted.
-            } else {
-              const sharedQueue: SharedQueue<any> = (updateQueue: any).shared;
-              const pending = sharedQueue.pending;
-              if (pending === null) {
-                // This is the first update. Create a circular list.
-                update.next = update;
+              // Inlined `enqueueUpdate` to remove interleaved update check
+              const updateQueue = fiber.updateQueue;
+              if (updateQueue === null) {
+                // Only occurs if the fiber has been unmounted.
               } else {
-                update.next = pending.next;
-                pending.next = update;
+                const sharedQueue: SharedQueue<any> = (updateQueue: any).shared;
+                const pending = sharedQueue.pending;
+                if (pending === null) {
+                  // This is the first update. Create a circular list.
+                  update.next = update;
+                } else {
+                  update.next = pending.next;
+                  pending.next = update;
+                }
+                sharedQueue.pending = update;
               }
-              sharedQueue.pending = update;
             }
+            // Mark the updated lanes on the list, too.
+            list.lanes = mergeLanes(list.lanes, renderLanes);
           }
+
           fiber.lanes = mergeLanes(fiber.lanes, renderLanes);
           const alternate = fiber.alternate;
           if (alternate !== null) {
@@ -244,9 +264,6 @@ export function propagateContextChange<T>(
           }
           scheduleWorkOnParentPath(fiber.return, renderLanes);
 
-          // Mark the updated lanes on the list, too.
-          list.lanes = mergeLanes(list.lanes, renderLanes);
-
           // Since we already found a match, we can stop traversing the
           // dependency list.
           break;
@@ -311,6 +328,30 @@ export function propagateContextChange<T>(
   }
 }
 
+export function checkIfContextChanged(currentDependencies: Dependencies) {
+  if (!enableLazyContextPropagation) {
+    return false;
+  }
+  // Iterate over the current dependencies to see if something changed. This
+  // only gets called if props and state has already bailed out, so it's a
+  // relatively uncommon path, except at the root of a changed subtree.
+  // Alternatively, we could move these comparisons into `readContext`, but
+  // that's a much hotter path, so I think this is an appropriate trade off.
+  let dependency = currentDependencies.firstContext;
+  while (dependency !== null) {
+    const context = dependency.context;
+    const newValue = isPrimaryRenderer
+      ? context._currentValue
+      : context._currentValue2;
+    const oldValue = dependency.memoizedValue;
+    if (!is(newValue, oldValue)) {
+      return true;
+    }
+    dependency = dependency.next;
+  }
+  return false;
+}
+
 export function prepareToReadContext(
   workInProgress: Fiber,
   renderLanes: Lanes,
@@ -321,14 +362,19 @@ export function prepareToReadContext(
 
   const dependencies = workInProgress.dependencies;
   if (dependencies !== null) {
-    const firstContext = dependencies.firstContext;
-    if (firstContext !== null) {
-      if (includesSomeLane(dependencies.lanes, renderLanes)) {
-        // Context list has a pending update. Mark that this fiber performed work.
-        markWorkInProgressReceivedUpdate();
-      }
+    if (enableLazyContextPropagation) {
       // Reset the work-in-progress list
       dependencies.firstContext = null;
+    } else {
+      const firstContext = dependencies.firstContext;
+      if (firstContext !== null) {
+        if (includesSomeLane(dependencies.lanes, renderLanes)) {
+          // Context list has a pending update. Mark that this fiber performed work.
+          markWorkInProgressReceivedUpdate();
+        }
+        // Reset the work-in-progress list
+        dependencies.firstContext = null;
+      }
     }
   }
 }
@@ -350,6 +396,10 @@ export function readContext<T>(
     }
   }
 
+  const value = isPrimaryRenderer
+    ? context._currentValue
+    : context._currentValue2;
+
   if (lastContextWithAllBitsObserved === context) {
     // Nothing to do. We already observe everything in this context.
   } else if (observedBits === false || observedBits === 0) {
@@ -370,6 +420,7 @@ export function readContext<T>(
     const contextItem = {
       context: ((context: any): ReactContext<mixed>),
       observedBits: resolvedObservedBits,
+      memoizedValue: value,
       next: null,
     };
 
@@ -387,6 +438,8 @@ export function readContext<T>(
       currentlyRenderingFiber.dependencies = {
         lanes: NoLanes,
         firstContext: contextItem,
+
+        // TODO: This is an old field. Delete it.
         responders: null,
       };
     } else {
@@ -394,5 +447,5 @@ export function readContext<T>(
       lastContextDependency = lastContextDependency.next = contextItem;
     }
   }
-  return isPrimaryRenderer ? context._currentValue : context._currentValue2;
+  return value;
 }

packages/react-reconciler/src/ReactInternalTypes.js

@@ -51,6 +51,7 @@ export type ContextDependency<T> = {
   context: ReactContext<T>,
   observedBits: number,
   next: ContextDependency<mixed> | null,
+  memoizedValue: T,
   ...
 };
 

packages/react-reconciler/src/__tests__/ReactContextPropagation-test.js

@@ -0,0 +1,206 @@
+let React;
+let ReactNoop;
+let Scheduler;
+let useState;
+let useContext;
+
+describe('ReactLazyContextPropagation', () => {
+  beforeEach(() => {
+    jest.resetModules();
+
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+    Scheduler = require('scheduler');
+    useState = React.useState;
+    useContext = React.useContext;
+  });
+
+  function Text({text}) {
+    Scheduler.unstable_yieldValue(text);
+    return text;
+  }
+
+  // NOTE: These tests are not specific to the lazy propagation (as opposed to
+  // eager propagation). The behavior should be the same in both
+  // implementations. These are tests that are more relevant to the lazy
+  // propagation implementation, though.
+
+  test(
+    'context change should prevent bailout of memoized component (useMemo -> ' +
+      'no intermediate fiber)',
+    async () => {
+      const root = ReactNoop.createRoot();
+
+      const Context = React.createContext(0);
+
+      let setValue;
+      function App() {
+        const [value, _setValue] = useState(0);
+        setValue = _setValue;
+
+        // NOTE: It's an important part of this test that we're memoizing the
+        // props of the Consumer component, as opposed to wrapping in an
+        // additional memoized fiber, because the implementation propagates
+        // context changes whenever a fiber bails out.
+        const consumer = React.useMemo(() => <Consumer />, []);
+
+        return <Context.Provider value={value}>{consumer}</Context.Provider>;
+      }
+
+      function Consumer() {
+        const value = useContext(Context);
+        // Even though Consumer is memoized, Consumer should re-render
+        // DeepChild whenever the context value changes. Otherwise DeepChild
+        // won't receive the new value.
+        return <DeepChild value={value} />;
+      }
+
+      function DeepChild({value}) {
+        return <Text text={value} />;
+      }
+
+      await ReactNoop.act(async () => {
+        root.render(<App />);
+      });
+      expect(Scheduler).toHaveYielded([0]);
+      expect(root).toMatchRenderedOutput('0');
+
+      await ReactNoop.act(async () => {
+        setValue(1);
+      });
+      expect(Scheduler).toHaveYielded([1]);
+      expect(root).toMatchRenderedOutput('1');
+    },
+  );
+
+  test('context change should prevent bailout of memoized component (memo HOC)', async () => {
+    const root = ReactNoop.createRoot();
+
+    const Context = React.createContext(0);
+
+    let setValue;
+    function App() {
+      const [value, _setValue] = useState(0);
+      setValue = _setValue;
+      return (
+        <Context.Provider value={value}>
+          <Consumer />
+        </Context.Provider>
+      );
+    }
+
+    const Consumer = React.memo(() => {
+      const value = useContext(Context);
+      // Even though Consumer is memoized, Consumer should re-render
+      // DeepChild whenever the context value changes. Otherwise DeepChild
+      // won't receive the new value.
+      return <DeepChild value={value} />;
+    });
+
+    function DeepChild({value}) {
+      return <Text text={value} />;
+    }
+
+    await ReactNoop.act(async () => {
+      root.render(<App />);
+    });
+    expect(Scheduler).toHaveYielded([0]);
+    expect(root).toMatchRenderedOutput('0');
+
+    await ReactNoop.act(async () => {
+      setValue(1);
+    });
+    expect(Scheduler).toHaveYielded([1]);
+    expect(root).toMatchRenderedOutput('1');
+  });
+
+  test('context change should prevent bailout of memoized component (PureComponent)', async () => {
+    const root = ReactNoop.createRoot();
+
+    const Context = React.createContext(0);
+
+    let setValue;
+    function App() {
+      const [value, _setValue] = useState(0);
+      setValue = _setValue;
+      return (
+        <Context.Provider value={value}>
+          <Consumer />
+        </Context.Provider>
+      );
+    }
+
+    class Consumer extends React.PureComponent {
+      static contextType = Context;
+      render() {
+        // Even though Consumer is memoized, Consumer should re-render
+        // DeepChild whenever the context value changes. Otherwise DeepChild
+        // won't receive the new value.
+        return <DeepChild value={this.context} />;
+      }
+    }
+
+    function DeepChild({value}) {
+      return <Text text={value} />;
+    }
+
+    await ReactNoop.act(async () => {
+      root.render(<App />);
+    });
+    expect(Scheduler).toHaveYielded([0]);
+    expect(root).toMatchRenderedOutput('0');
+
+    await ReactNoop.act(async () => {
+      setValue(1);
+    });
+    expect(Scheduler).toHaveYielded([1]);
+    expect(root).toMatchRenderedOutput('1');
+  });
+
+  test("context consumer bails out if context hasn't changed", async () => {
+    const root = ReactNoop.createRoot();
+
+    const Context = React.createContext(0);
+
+    function App() {
+      return (
+        <Context.Provider value={0}>
+          <Consumer />
+        </Context.Provider>
+      );
+    }
+
+    let setOtherValue;
+    const Consumer = React.memo(() => {
+      const value = useContext(Context);
+
+      const [, _setOtherValue] = useState(0);
+      setOtherValue = _setOtherValue;
+
+      Scheduler.unstable_yieldValue('Consumer');
+
+      return <Text text={value} />;
+    });
+
+    await ReactNoop.act(async () => {
+      root.render(<App />);
+    });
+    expect(Scheduler).toHaveYielded(['Consumer', 0]);
+    expect(root).toMatchRenderedOutput('0');
+
+    await ReactNoop.act(async () => {
+      // Intentionally calling setState to some other arbitrary value before
+      // setting it back to the current one. That way an update is scheduled,
+      // but we'll bail out during render when nothing has changed.
+      setOtherValue(1);
+      setOtherValue(0);
+    });
+    // NOTE: If this didn't yield anything, that indicates that we never visited
+    // the consumer during the render phase, which probably means the eager
+    // bailout mechanism kicked in. Because we're testing the _lazy_ bailout
+    // mechanism, update this test to foil the _eager_ bailout, somehow. Perhaps
+    // by switching to useReducer.
+    expect(Scheduler).toHaveYielded(['Consumer']);
+    expect(root).toMatchRenderedOutput('0');
+  });
+});

packages/react/src/__tests__/ReactContextValidator-test.js

@@ -431,8 +431,12 @@ describe('ReactContextValidator', () => {
     expect(renderContext).toBe(secondContext);
     expect(componentDidUpdateContext).toBe(secondContext);
 
-    // sCU is not called in this case because React force updates when a provider re-renders
-    expect(shouldComponentUpdateWasCalled).toBe(false);
+    if (gate(flags => flags.enableLazyContextPropagation)) {
+      expect(shouldComponentUpdateWasCalled).toBe(true);
+    } else {
+      // sCU is not called in this case because React force updates when a provider re-renders
+      expect(shouldComponentUpdateWasCalled).toBe(false);
+    }
   });
 
   it('should re-render PureComponents when context Provider updates', () => {

packages/shared/ReactFeatureFlags.js

@@ -158,3 +158,5 @@ export const enableDiscreteEventMicroTasks = false;
 export const enableSyncMicroTasks = false;
 
 export const enableNativeEventPriorityInference = false;
+
+export const enableLazyContextPropagation = false;

packages/shared/forks/ReactFeatureFlags.native-fb.js

@@ -61,6 +61,7 @@ export const disableSchedulerTimeoutInWorkLoop = false;
 export const enableDiscreteEventMicroTasks = false;
 export const enableSyncMicroTasks = false;
 export const enableNativeEventPriorityInference = false;
+export const enableLazyContextPropagation = false;
 
 // Flow magic to verify the exports of this file match the original version.
 // eslint-disable-next-line no-unused-vars

packages/shared/forks/ReactFeatureFlags.native-oss.js

@@ -60,6 +60,7 @@ export const disableSchedulerTimeoutInWorkLoop = false;
 export const enableDiscreteEventMicroTasks = false;
 export const enableSyncMicroTasks = false;
 export const enableNativeEventPriorityInference = false;
+export const enableLazyContextPropagation = false;
 
 // Flow magic to verify the exports of this file match the original version.
 // eslint-disable-next-line no-unused-vars

packages/shared/forks/ReactFeatureFlags.test-renderer.js

@@ -60,6 +60,7 @@ export const disableSchedulerTimeoutInWorkLoop = false;
 export const enableDiscreteEventMicroTasks = false;
 export const enableSyncMicroTasks = false;
 export const enableNativeEventPriorityInference = false;
+export const enableLazyContextPropagation = false;
 
 // Flow magic to verify the exports of this file match the original version.
 // eslint-disable-next-line no-unused-vars

packages/shared/forks/ReactFeatureFlags.test-renderer.native.js

@@ -60,6 +60,7 @@ export const disableSchedulerTimeoutInWorkLoop = false;
 export const enableDiscreteEventMicroTasks = false;
 export const enableSyncMicroTasks = false;
 export const enableNativeEventPriorityInference = false;
+export const enableLazyContextPropagation = false;
 
 // Flow magic to verify the exports of this file match the original version.
 // eslint-disable-next-line no-unused-vars

packages/shared/forks/ReactFeatureFlags.test-renderer.www.js

@@ -60,6 +60,7 @@ export const disableSchedulerTimeoutInWorkLoop = false;
 export const enableDiscreteEventMicroTasks = false;
 export const enableSyncMicroTasks = false;
 export const enableNativeEventPriorityInference = false;
+export const enableLazyContextPropagation = false;
 
 // Flow magic to verify the exports of this file match the original version.
 // eslint-disable-next-line no-unused-vars

packages/shared/forks/ReactFeatureFlags.testing.js

@@ -60,6 +60,7 @@ export const disableSchedulerTimeoutInWorkLoop = false;
 export const enableDiscreteEventMicroTasks = false;
 export const enableSyncMicroTasks = false;
 export const enableNativeEventPriorityInference = false;
+export const enableLazyContextPropagation = false;
 
 // Flow magic to verify the exports of this file match the original version.
 // eslint-disable-next-line no-unused-vars

packages/shared/forks/ReactFeatureFlags.testing.www.js

@@ -60,6 +60,7 @@ export const disableSchedulerTimeoutInWorkLoop = false;
 export const enableDiscreteEventMicroTasks = false;
 export const enableSyncMicroTasks = false;
 export const enableNativeEventPriorityInference = false;
+export const enableLazyContextPropagation = false;
 
 // Flow magic to verify the exports of this file match the original version.
 // eslint-disable-next-line no-unused-vars

packages/shared/forks/ReactFeatureFlags.www-dynamic.js

@@ -59,3 +59,4 @@ export const disableSchedulerTimeoutInWorkLoop = __VARIANT__;
 export const enableDiscreteEventMicroTasks = __VARIANT__;
 export const enableSyncMicroTasks = __VARIANT__;
 export const enableNativeEventPriorityInference = __VARIANT__;
+export const enableLazyContextPropagation = __VARIANT__;

packages/shared/forks/ReactFeatureFlags.www.js

@@ -35,6 +35,7 @@ export const {
   enableDiscreteEventMicroTasks,
   enableSyncMicroTasks,
   enableNativeEventPriorityInference,
+  enableLazyContextPropagation,
 } = dynamicFeatureFlags;
 
 // On WWW, __EXPERIMENTAL__ is used for a new modern build.