Fork readContextForConsumer

Author: Jack Pope

packages/react-debug-tools/src/ReactDebugHooks.js

@@ -37,6 +37,7 @@ import {
   REACT_CONTEXT_TYPE,
 } from 'shared/ReactSymbols';
 import hasOwnProperty from 'shared/hasOwnProperty';
+import type {ContextDependencyWithCompare} from '../../react-reconciler/src/ReactInternalTypes';
 
 type CurrentDispatcherRef = typeof ReactSharedInternals;
 
@@ -155,7 +156,10 @@ function getPrimitiveStackCache(): Map<string, Array<any>> {
 
 let currentFiber: null | Fiber = null;
 let currentHook: null | Hook = null;
-let currentContextDependency: null | ContextDependency<mixed, mixed> = null;
+let currentContextDependency:
+  | null
+  | ContextDependency<mixed>
+  | ContextDependencyWithCompare<mixed, mixed> = null;
 
 function nextHook(): null | Hook {
   const hook = currentHook;

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1060,7 +1060,7 @@ function updateWorkInProgressHook(): Hook {
 
 function unstable_useContextWithBailout<T>(
   context: ReactContext<T>,
-  compare: void | (T => mixed),
+  compare: (T => mixed) | null,
 ): T {
   return readContextAndCompare(context, compare);
 }
@@ -4049,7 +4049,7 @@ if (__DEV__) {
   }
   if (enableContextProfiling) {
     (HooksDispatcherOnMountInDEV: Dispatcher).unstable_useContextWithBailout =
-      function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
+      function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
         currentHookNameInDev = 'useContext';
         mountHookTypesDev();
         return unstable_useContextWithBailout(context, compare);
@@ -4238,7 +4238,7 @@ if (__DEV__) {
   }
   if (enableContextProfiling) {
     (HooksDispatcherOnMountWithHookTypesInDEV: Dispatcher).unstable_useContextWithBailout =
-      function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
+      function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
         currentHookNameInDev = 'useContext';
         updateHookTypesDev();
         return unstable_useContextWithBailout(context, compare);
@@ -4426,7 +4426,7 @@ if (__DEV__) {
   }
   if (enableContextProfiling) {
     (HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
-      function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
+      function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
         currentHookNameInDev = 'useContext';
         updateHookTypesDev();
         return unstable_useContextWithBailout(context, compare);
@@ -4614,7 +4614,7 @@ if (__DEV__) {
   }
   if (enableContextProfiling) {
     (HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
-      function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
+      function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
         currentHookNameInDev = 'useContext';
         updateHookTypesDev();
         return unstable_useContextWithBailout(context, compare);
@@ -4828,7 +4828,7 @@ if (__DEV__) {
   }
   if (enableContextProfiling) {
     (HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
-      function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
+      function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
         currentHookNameInDev = 'useContext';
         warnInvalidHookAccess();
         mountHookTypesDev();
@@ -5043,7 +5043,7 @@ if (__DEV__) {
   }
   if (enableContextProfiling) {
     (InvalidNestedHooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
-      function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
+      function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
         currentHookNameInDev = 'useContext';
         warnInvalidHookAccess();
         updateHookTypesDev();
@@ -5258,7 +5258,7 @@ if (__DEV__) {
   }
   if (enableContextProfiling) {
     (InvalidNestedHooksDispatcherOnRerenderInDEV: Dispatcher).unstable_useContextWithBailout =
-      function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
+      function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
         currentHookNameInDev = 'useContext';
         warnInvalidHookAccess();
         updateHookTypesDev();

packages/react-reconciler/src/ReactFiberNewContext.js

@@ -12,6 +12,7 @@ import type {
   Fiber,
   ContextDependency,
   Dependencies,
+  ContextDependencyWithCompare,
 } from './ReactInternalTypes';
 import type {StackCursor} from './ReactFiberStack';
 import type {Lanes} from './ReactFiberLane';
@@ -72,7 +73,10 @@ if (__DEV__) {
 }
 
 let currentlyRenderingFiber: Fiber | null = null;
-let lastContextDependency: ContextDependency<mixed, mixed> | null = null;
+let lastContextDependency:
+  | ContextDependency<mixed>
+  | ContextDependencyWithCompare<mixed, mixed>
+  | null = null;
 let lastFullyObservedContext: ReactContext<any> | null = null;
 
 let isDisallowedContextReadInDEV: boolean = false;
@@ -403,19 +407,21 @@ function propagateContextChanges<T>(
           const context: ReactContext<T> = contexts[i];
           // Check if the context matches.
           if (dependency.context === context) {
-            const compare = dependency.compare;
-            if (enableContextProfiling && compare != null) {
-              const newValue = isPrimaryRenderer
-                ? dependency.context._currentValue
-                : dependency.context._currentValue2;
-              if (
-                !checkIfComparedContextValuesChanged(
-                  dependency.lastComparedValue,
-                  compare(newValue),
-                )
-              ) {
-                // Compared value hasn't changed. Bail out early.
-                continue findContext;
+            if (enableContextProfiling) {
+              const compare = dependency.compare;
+              if (compare != null) {
+                const newValue = isPrimaryRenderer
+                  ? dependency.context._currentValue
+                  : dependency.context._currentValue2;
+                if (
+                  !checkIfComparedContextValuesChanged(
+                    dependency.lastComparedValue,
+                    compare(newValue),
+                  )
+                ) {
+                  // Compared value hasn't changed. Bail out early.
+                  continue findContext;
+                }
               }
             }
             // Match! Schedule an update on this fiber.
@@ -697,12 +703,11 @@ export function checkIfContextChanged(
       ? context._currentValue
       : context._currentValue2;
     const oldValue = dependency.memoizedValue;
-    const compare = dependency.compare;
-    if (enableContextProfiling && compare != null) {
+    if (enableContextProfiling && dependency.compare != null) {
       if (
         checkIfComparedContextValuesChanged(
           dependency.lastComparedValue,
-          compare(newValue),
+          dependency.compare(newValue),
         )
       ) {
         return true;
@@ -746,13 +751,17 @@ export function prepareToReadContext(
 
 export function readContextAndCompare<C>(
   context: ReactContext<C>,
-  compare: void | (C => mixed),
+  compare: (C => mixed) | null,
 ): C {
   if (!enableLazyContextPropagation) {
     return readContext(context);
   }
 
-  return readContextForConsumer(currentlyRenderingFiber, context, compare);
+  return readContextForConsumer_withCompare(
+    currentlyRenderingFiber,
+    context,
+    compare,
+  );
 }
 
 export function readContext<T>(context: ReactContext<T>): T {
@@ -782,12 +791,12 @@ export function readContextDuringReconciliation<T>(
   return readContextForConsumer(consumer, context);
 }
 
-type ContextCompare<C, S> = C => S;
+type ContextCompare<C, V> = C => V | null;
 
-function readContextForConsumer<C, S>(
+function readContextForConsumer_withCompare<C, S>(
   consumer: Fiber | null,
   context: ReactContext<C>,
-  compare?: void | (C => S),
+  compare: (C => S) | null,
 ): C {
   const value = isPrimaryRenderer
     ? context._currentValue
@@ -800,7 +809,7 @@ function readContextForConsumer<C, S>(
       context: ((context: any): ReactContext<mixed>),
       memoizedValue: value,
       next: null,
-      compare: ((compare: any): ContextCompare<mixed, mixed> | null),
+      compare: compare ? ((compare: any): ContextCompare<mixed, mixed>) : null,
       lastComparedValue: compare != null ? compare(value) : null,
     };
 
@@ -830,3 +839,47 @@ function readContextForConsumer<C, S>(
   }
   return value;
 }
+
+function readContextForConsumer<C>(
+  consumer: Fiber | null,
+  context: ReactContext<C>,
+): C {
+  const value = isPrimaryRenderer
+    ? context._currentValue
+    : context._currentValue2;
+
+  if (lastFullyObservedContext === context) {
+    // Nothing to do. We already observe everything in this context.
+  } else {
+    const contextItem = {
+      context: ((context: any): ReactContext<mixed>),
+      memoizedValue: value,
+      next: null,
+    };
+
+    if (lastContextDependency === null) {
+      if (consumer === null) {
+        throw new Error(
+          'Context can only be read while React is rendering. ' +
+            'In classes, you can read it in the render method or getDerivedStateFromProps. ' +
+            'In function components, you can read it directly in the function body, but not ' +
+            'inside Hooks like useReducer() or useMemo().',
+        );
+      }
+
+      // This is the first dependency for this component. Create a new list.
+      lastContextDependency = contextItem;
+      consumer.dependencies = {
+        lanes: NoLanes,
+        firstContext: contextItem,
+      };
+      if (enableLazyContextPropagation) {
+        consumer.flags |= NeedsPropagation;
+      }
+    } else {
+      // Append a new context item.
+      lastContextDependency = lastContextDependency.next = contextItem;
+    }
+  }
+  return value;
+}

packages/react-reconciler/src/ReactInternalTypes.js

@@ -61,18 +61,32 @@ export type HookType =
   | 'useFormState'
   | 'useActionState';
 
-export type ContextDependency<C, S> = {
+export type ContextDependency<C> = {
   context: ReactContext<C>,
-  next: ContextDependency<mixed, mixed> | null,
+  next:
+    | ContextDependency<mixed>
+    | ContextDependencyWithCompare<mixed, mixed>
+    | null,
+  memoizedValue: C,
+};
+
+export type ContextDependencyWithCompare<C, S> = {
+  context: ReactContext<C>,
+  next:
+    | ContextDependency<mixed>
+    | ContextDependencyWithCompare<mixed, mixed>
+    | null,
   memoizedValue: C,
   compare: (C => S) | null,
-  lastComparedValue: S | null,
-  ...
+  lastComparedValue?: S | null,
 };
 
 export type Dependencies = {
   lanes: Lanes,
-  firstContext: ContextDependency<mixed, mixed> | null,
+  firstContext:
+    | ContextDependency<mixed>
+    | ContextDependencyWithCompare<mixed, mixed>
+    | null,
   ...
 };
 
@@ -388,7 +402,7 @@ export type Dispatcher = {
   ): [S, Dispatch<A>],
   unstable_useContextWithBailout?: <T>(
     context: ReactContext<T>,
-    compare: void | (T => mixed),
+    compare: (T => mixed) | null,
   ) => T,
   useContext<T>(context: ReactContext<T>): T,
   useRef<T>(initialValue: T): {current: T},

packages/react/src/ReactHooks.js

@@ -71,7 +71,7 @@ export function useContext<T>(Context: ReactContext<T>): T {
 
 export function unstable_useContextWithBailout<T>(
   context: ReactContext<T>,
-  compare: void | (T => mixed),
+  compare: (T => mixed) | null,
 ): T {
   if (!(enableLazyContextPropagation && enableContextProfiling)) {
     throw new Error('Not implemented.');