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;
 }