Add unstable context bailout for profiling (#30407)

**This API is not intended to ship. This is a temporary unstable hook
for internal performance profiling.**

This PR exposes `unstable_useContextWithBailout`, which takes a compare
function in addition to Context. The comparison function is run to
determine if Context propagation and render should bail out earlier.
`unstable_useContextWithBailout` returns the full Context value, same as
`useContext`.

We can profile this API against `useContext` to better measure the cost
of Context value updates and gather more data around propagation and
render performance.

The bailout logic and test cases are based on
https://github.com/facebook/react/pull/20646

Additionally, this implementation allows multiple values to be compared
in one hook by returning a tuple to avoid requiring additional Context
consumer hooks.

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 {ContextDependencyWithSelect} 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> = null;
+let currentContextDependency:
+  | null
+  | ContextDependency<mixed>
+  | ContextDependencyWithSelect<mixed> = null;
 
 function nextHook(): null | Hook {
   const hook = currentHook;

packages/react-reconciler/src/ReactFiberHooks.js

@@ -47,6 +47,7 @@ import {
   enableUseDeferredValueInitialArg,
   disableLegacyMode,
   enableNoCloningMemoCache,
+  enableContextProfiling,
 } from 'shared/ReactFeatureFlags';
 import {
   REACT_CONTEXT_TYPE,
@@ -81,7 +82,11 @@ import {
   ContinuousEventPriority,
   higherEventPriority,
 } from './ReactEventPriorities';
-import {readContext, checkIfContextChanged} from './ReactFiberNewContext';
+import {
+  readContext,
+  readContextAndCompare,
+  checkIfContextChanged,
+} from './ReactFiberNewContext';
 import {HostRoot, CacheComponent, HostComponent} from './ReactWorkTags';
 import {
   LayoutStatic as LayoutStaticEffect,
@@ -1053,6 +1058,16 @@ function updateWorkInProgressHook(): Hook {
   return workInProgressHook;
 }
 
+function unstable_useContextWithBailout<T>(
+  context: ReactContext<T>,
+  select: (T => Array<mixed>) | null,
+): T {
+  if (select === null) {
+    return readContext(context);
+  }
+  return readContextAndCompare(context, select);
+}
+
 // NOTE: defining two versions of this function to avoid size impact when this feature is disabled.
 // Previously this function was inlined, the additional `memoCache` property makes it not inlined.
 let createFunctionComponentUpdateQueue: () => FunctionComponentUpdateQueue;
@@ -3689,6 +3704,10 @@ if (enableAsyncActions) {
 if (enableAsyncActions) {
   (ContextOnlyDispatcher: Dispatcher).useOptimistic = throwInvalidHookError;
 }
+if (enableContextProfiling) {
+  (ContextOnlyDispatcher: Dispatcher).unstable_useContextWithBailout =
+    throwInvalidHookError;
+}
 
 const HooksDispatcherOnMount: Dispatcher = {
   readContext,
@@ -3728,6 +3747,10 @@ if (enableAsyncActions) {
 if (enableAsyncActions) {
   (HooksDispatcherOnMount: Dispatcher).useOptimistic = mountOptimistic;
 }
+if (enableContextProfiling) {
+  (HooksDispatcherOnMount: Dispatcher).unstable_useContextWithBailout =
+    unstable_useContextWithBailout;
+}
 
 const HooksDispatcherOnUpdate: Dispatcher = {
   readContext,
@@ -3767,6 +3790,10 @@ if (enableAsyncActions) {
 if (enableAsyncActions) {
   (HooksDispatcherOnUpdate: Dispatcher).useOptimistic = updateOptimistic;
 }
+if (enableContextProfiling) {
+  (HooksDispatcherOnUpdate: Dispatcher).unstable_useContextWithBailout =
+    unstable_useContextWithBailout;
+}
 
 const HooksDispatcherOnRerender: Dispatcher = {
   readContext,
@@ -3806,6 +3833,10 @@ if (enableAsyncActions) {
 if (enableAsyncActions) {
   (HooksDispatcherOnRerender: Dispatcher).useOptimistic = rerenderOptimistic;
 }
+if (enableContextProfiling) {
+  (HooksDispatcherOnRerender: Dispatcher).unstable_useContextWithBailout =
+    unstable_useContextWithBailout;
+}
 
 let HooksDispatcherOnMountInDEV: Dispatcher | null = null;
 let HooksDispatcherOnMountWithHookTypesInDEV: Dispatcher | null = null;
@@ -4019,6 +4050,17 @@ if (__DEV__) {
         return mountOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnMountInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        mountHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   HooksDispatcherOnMountWithHookTypesInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4200,6 +4242,17 @@ if (__DEV__) {
         return mountOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnMountWithHookTypesInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   HooksDispatcherOnUpdateInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4380,6 +4433,17 @@ if (__DEV__) {
         return updateOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   HooksDispatcherOnRerenderInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4560,6 +4624,17 @@ if (__DEV__) {
         return rerenderOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   InvalidNestedHooksDispatcherOnMountInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4766,6 +4841,18 @@ if (__DEV__) {
         return mountOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        warnInvalidHookAccess();
+        mountHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   InvalidNestedHooksDispatcherOnUpdateInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -4972,6 +5059,18 @@ if (__DEV__) {
         return updateOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (InvalidNestedHooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        warnInvalidHookAccess();
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 
   InvalidNestedHooksDispatcherOnRerenderInDEV = {
     readContext<T>(context: ReactContext<T>): T {
@@ -5178,4 +5277,16 @@ if (__DEV__) {
         return rerenderOptimistic(passthrough, reducer);
       };
   }
+  if (enableContextProfiling) {
+    (InvalidNestedHooksDispatcherOnRerenderInDEV: Dispatcher).unstable_useContextWithBailout =
+      function <T>(
+        context: ReactContext<T>,
+        select: (T => Array<mixed>) | null,
+      ): T {
+        currentHookNameInDev = 'useContext';
+        warnInvalidHookAccess();
+        updateHookTypesDev();
+        return unstable_useContextWithBailout(context, select);
+      };
+  }
 }

packages/react-reconciler/src/ReactFiberNewContext.js

@@ -12,6 +12,7 @@ import type {
   Fiber,
   ContextDependency,
   Dependencies,
+  ContextDependencyWithSelect,
 } from './ReactInternalTypes';
 import type {StackCursor} from './ReactFiberStack';
 import type {Lanes} from './ReactFiberLane';
@@ -51,6 +52,8 @@ import {
   getHostTransitionProvider,
   HostTransitionContext,
 } from './ReactFiberHostContext';
+import isArray from '../../shared/isArray';
+import {enableContextProfiling} from '../../shared/ReactFeatureFlags';
 
 const valueCursor: StackCursor<mixed> = createCursor(null);
 
@@ -70,7 +73,10 @@ if (__DEV__) {
 }
 
 let currentlyRenderingFiber: Fiber | null = null;
-let lastContextDependency: ContextDependency<mixed> | null = null;
+let lastContextDependency:
+  | ContextDependency<mixed>
+  | ContextDependencyWithSelect<mixed>
+  | null = null;
 let lastFullyObservedContext: ReactContext<any> | null = null;
 
 let isDisallowedContextReadInDEV: boolean = false;
@@ -400,8 +406,24 @@ function propagateContextChanges<T>(
         findContext: for (let i = 0; i < contexts.length; i++) {
           const context: ReactContext<T> = contexts[i];
           // Check if the context matches.
-          // TODO: Compare selected values to bail out early.
           if (dependency.context === context) {
+            if (enableContextProfiling) {
+              const select = dependency.select;
+              if (select != null && dependency.lastSelectedValue != null) {
+                const newValue = isPrimaryRenderer
+                  ? dependency.context._currentValue
+                  : dependency.context._currentValue2;
+                if (
+                  !checkIfSelectedContextValuesChanged(
+                    dependency.lastSelectedValue,
+                    select(newValue),
+                  )
+                ) {
+                  // Compared value hasn't changed. Bail out early.
+                  continue findContext;
+                }
+              }
+            }
             // Match! Schedule an update on this fiber.
 
             // In the lazy implementation, don't mark a dirty flag on the
@@ -641,6 +663,29 @@ function propagateParentContextChanges(
   workInProgress.flags |= DidPropagateContext;
 }
 
+function checkIfSelectedContextValuesChanged(
+  oldComparedValue: Array<mixed>,
+  newComparedValue: Array<mixed>,
+): boolean {
+  // We have an implicit contract that compare functions must return arrays.
+  // This allows us to compare multiple values in the same context access
+  // since compiling to additional hook calls regresses perf.
+  if (isArray(oldComparedValue) && isArray(newComparedValue)) {
+    if (oldComparedValue.length !== newComparedValue.length) {
+      return true;
+    }
+
+    for (let i = 0; i < oldComparedValue.length; i++) {
+      if (!is(newComparedValue[i], oldComparedValue[i])) {
+        return true;
+      }
+    }
+  } else {
+    throw new Error('Compared context values must be arrays');
+  }
+  return false;
+}
+
 export function checkIfContextChanged(
   currentDependencies: Dependencies,
 ): boolean {
@@ -659,8 +704,23 @@ export function checkIfContextChanged(
       ? context._currentValue
       : context._currentValue2;
     const oldValue = dependency.memoizedValue;
-    if (!is(newValue, oldValue)) {
-      return true;
+    if (
+      enableContextProfiling &&
+      dependency.select != null &&
+      dependency.lastSelectedValue != null
+    ) {
+      if (
+        checkIfSelectedContextValuesChanged(
+          dependency.lastSelectedValue,
+          dependency.select(newValue),
+        )
+      ) {
+        return true;
+      }
+    } else {
+      if (!is(newValue, oldValue)) {
+        return true;
+      }
     }
     dependency = dependency.next;
   }
@@ -694,6 +754,21 @@ export function prepareToReadContext(
   }
 }
 
+export function readContextAndCompare<C>(
+  context: ReactContext<C>,
+  select: C => Array<mixed>,
+): C {
+  if (!(enableLazyContextPropagation && enableContextProfiling)) {
+    throw new Error('Not implemented.');
+  }
+
+  return readContextForConsumer_withSelect(
+    currentlyRenderingFiber,
+    context,
+    select,
+  );
+}
+
 export function readContext<T>(context: ReactContext<T>): T {
   if (__DEV__) {
     // This warning would fire if you read context inside a Hook like useMemo.
@@ -721,10 +796,57 @@ export function readContextDuringReconciliation<T>(
   return readContextForConsumer(consumer, context);
 }
 
-function readContextForConsumer<T>(
+function readContextForConsumer_withSelect<C>(
   consumer: Fiber | null,
-  context: ReactContext<T>,
-): T {
+  context: ReactContext<C>,
+  select: C => Array<mixed>,
+): 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,
+      select: ((select: any): (context: mixed) => Array<mixed>),
+      lastSelectedValue: select(value),
+    };
+
+    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;
+}
+
+function readContextForConsumer<C>(
+  consumer: Fiber | null,
+  context: ReactContext<C>,
+): C {
   const value = isPrimaryRenderer
     ? context._currentValue
     : context._currentValue2;

packages/react-reconciler/src/ReactInternalTypes.js

@@ -61,16 +61,26 @@ export type HookType =
   | 'useFormState'
   | 'useActionState';
 
-export type ContextDependency<T> = {
-  context: ReactContext<T>,
-  next: ContextDependency<mixed> | null,
-  memoizedValue: T,
-  ...
+export type ContextDependency<C> = {
+  context: ReactContext<C>,
+  next: ContextDependency<mixed> | ContextDependencyWithSelect<mixed> | null,
+  memoizedValue: C,
+};
+
+export type ContextDependencyWithSelect<C> = {
+  context: ReactContext<C>,
+  next: ContextDependency<mixed> | ContextDependencyWithSelect<mixed> | null,
+  memoizedValue: C,
+  select: C => Array<mixed>,
+  lastSelectedValue: ?Array<mixed>,
 };
 
 export type Dependencies = {
   lanes: Lanes,
-  firstContext: ContextDependency<mixed> | null,
+  firstContext:
+    | ContextDependency<mixed>
+    | ContextDependencyWithSelect<mixed>
+    | null,
   ...
 };
 
@@ -384,6 +394,10 @@ export type Dispatcher = {
     initialArg: I,
     init?: (I) => S,
   ): [S, Dispatch<A>],
+  unstable_useContextWithBailout?: <T>(
+    context: ReactContext<T>,
+    select: (T => Array<mixed>) | null,
+  ) => T,
   useContext<T>(context: ReactContext<T>): T,
   useRef<T>(initialValue: T): {current: T},
   useEffect(

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

@@ -0,0 +1,217 @@
+let React;
+let ReactNoop;
+let Scheduler;
+let act;
+let assertLog;
+let useState;
+let useContext;
+let unstable_useContextWithBailout;
+
+describe('ReactContextWithBailout', () => {
+  beforeEach(() => {
+    jest.resetModules();
+
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+    Scheduler = require('scheduler');
+    const testUtils = require('internal-test-utils');
+    act = testUtils.act;
+    assertLog = testUtils.assertLog;
+    useState = React.useState;
+    useContext = React.useContext;
+    unstable_useContextWithBailout = React.unstable_useContextWithBailout;
+  });
+
+  function Text({text}) {
+    Scheduler.log(text);
+    return text;
+  }
+
+  // @gate enableLazyContextPropagation && enableContextProfiling
+  test('unstable_useContextWithBailout basic usage', async () => {
+    const Context = React.createContext();
+
+    let setContext;
+    function App() {
+      const [context, _setContext] = useState({a: 'A0', b: 'B0', c: 'C0'});
+      setContext = _setContext;
+      return (
+        <Context.Provider value={context}>
+          <Indirection />
+        </Context.Provider>
+      );
+    }
+
+    // Intermediate parent that bails out. Children will only re-render when the
+    // context changes.
+    const Indirection = React.memo(() => {
+      return (
+        <>
+          A: <A />, B: <B />, C: <C />, AB: <AB />
+        </>
+      );
+    });
+
+    function A() {
+      const {a} = unstable_useContextWithBailout(Context, context => [
+        context.a,
+      ]);
+      return <Text text={a} />;
+    }
+
+    function B() {
+      const {b} = unstable_useContextWithBailout(Context, context => [
+        context.b,
+      ]);
+      return <Text text={b} />;
+    }
+
+    function C() {
+      const {c} = unstable_useContextWithBailout(Context, context => [
+        context.c,
+      ]);
+      return <Text text={c} />;
+    }
+
+    function AB() {
+      const {a, b} = unstable_useContextWithBailout(Context, context => [
+        context.a,
+        context.b,
+      ]);
+      return <Text text={a + b} />;
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      root.render(<App />);
+    });
+    assertLog(['A0', 'B0', 'C0', 'A0B0']);
+    expect(root).toMatchRenderedOutput('A: A0, B: B0, C: C0, AB: A0B0');
+
+    // Update a. Only the A and AB consumer should re-render.
+    await act(async () => {
+      setContext({a: 'A1', c: 'C0', b: 'B0'});
+    });
+    assertLog(['A1', 'A1B0']);
+    expect(root).toMatchRenderedOutput('A: A1, B: B0, C: C0, AB: A1B0');
+
+    // Update b. Only the B and AB consumer should re-render.
+    await act(async () => {
+      setContext({a: 'A1', b: 'B1', c: 'C0'});
+    });
+    assertLog(['B1', 'A1B1']);
+    expect(root).toMatchRenderedOutput('A: A1, B: B1, C: C0, AB: A1B1');
+
+    // Update c. Only the C consumer should re-render.
+    await act(async () => {
+      setContext({a: 'A1', b: 'B1', c: 'C1'});
+    });
+    assertLog(['C1']);
+    expect(root).toMatchRenderedOutput('A: A1, B: B1, C: C1, AB: A1B1');
+  });
+
+  // @gate enableLazyContextPropagation && enableContextProfiling
+  test('unstable_useContextWithBailout and useContext subscribing to same context in same component', async () => {
+    const Context = React.createContext();
+
+    let setContext;
+    function App() {
+      const [context, _setContext] = useState({a: 0, b: 0, unrelated: 0});
+      setContext = _setContext;
+      return (
+        <Context.Provider value={context}>
+          <Indirection />
+        </Context.Provider>
+      );
+    }
+
+    // Intermediate parent that bails out. Children will only re-render when the
+    // context changes.
+    const Indirection = React.memo(() => {
+      return <Child />;
+    });
+
+    function Child() {
+      const {a} = unstable_useContextWithBailout(Context, context => [
+        context.a,
+      ]);
+      const context = useContext(Context);
+      return <Text text={`A: ${a}, B: ${context.b}`} />;
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      root.render(<App />);
+    });
+    assertLog(['A: 0, B: 0']);
+    expect(root).toMatchRenderedOutput('A: 0, B: 0');
+
+    // Update an unrelated field that isn't used by the component. The context
+    // attempts to bail out, but the normal context forces an update.
+    await act(async () => {
+      setContext({a: 0, b: 0, unrelated: 1});
+    });
+    assertLog(['A: 0, B: 0']);
+    expect(root).toMatchRenderedOutput('A: 0, B: 0');
+  });
+
+  // @gate enableLazyContextPropagation && enableContextProfiling
+  test('unstable_useContextWithBailout and useContext subscribing to different contexts in same component', async () => {
+    const ContextA = React.createContext();
+    const ContextB = React.createContext();
+
+    let setContextA;
+    let setContextB;
+    function App() {
+      const [a, _setContextA] = useState({a: 0, unrelated: 0});
+      const [b, _setContextB] = useState(0);
+      setContextA = _setContextA;
+      setContextB = _setContextB;
+      return (
+        <ContextA.Provider value={a}>
+          <ContextB.Provider value={b}>
+            <Indirection />
+          </ContextB.Provider>
+        </ContextA.Provider>
+      );
+    }
+
+    // Intermediate parent that bails out. Children will only re-render when the
+    // context changes.
+    const Indirection = React.memo(() => {
+      return <Child />;
+    });
+
+    function Child() {
+      const {a} = unstable_useContextWithBailout(ContextA, context => [
+        context.a,
+      ]);
+      const b = useContext(ContextB);
+      return <Text text={`A: ${a}, B: ${b}`} />;
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      root.render(<App />);
+    });
+    assertLog(['A: 0, B: 0']);
+    expect(root).toMatchRenderedOutput('A: 0, B: 0');
+
+    // Update a field in A that isn't part of the compared context. It should
+    // bail out.
+    await act(async () => {
+      setContextA({a: 0, unrelated: 1});
+    });
+    assertLog([]);
+    expect(root).toMatchRenderedOutput('A: 0, B: 0');
+
+    // Now update the same a field again, but this time, also update a different
+    // context in the same batch. The other context prevents a bail out.
+    await act(async () => {
+      setContextA({a: 0, unrelated: 1});
+      setContextB(1);
+    });
+    assertLog(['A: 0, B: 1']);
+    expect(root).toMatchRenderedOutput('A: 0, B: 1');
+  });
+});

packages/react/index.fb.js

@@ -39,6 +39,7 @@ export {
   use,
   useActionState,
   useCallback,
+  unstable_useContextWithBailout,
   useContext,
   useDebugValue,
   useDeferredValue,

packages/react/src/ReactClient.js

@@ -38,6 +38,7 @@ import {postpone} from './ReactPostpone';
 import {
   getCacheForType,
   useCallback,
+  unstable_useContextWithBailout,
   useContext,
   useEffect,
   useEffectEvent,
@@ -83,6 +84,7 @@ export {
   cache,
   postpone as unstable_postpone,
   useCallback,
+  unstable_useContextWithBailout,
   useContext,
   useEffect,
   useEffectEvent as experimental_useEffectEvent,

packages/react/src/ReactHooks.js

@@ -19,6 +19,10 @@ import {REACT_CONSUMER_TYPE} from 'shared/ReactSymbols';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 
 import {enableAsyncActions} from 'shared/ReactFeatureFlags';
+import {
+  enableContextProfiling,
+  enableLazyContextPropagation,
+} from '../../shared/ReactFeatureFlags';
 
 type BasicStateAction<S> = (S => S) | S;
 type Dispatch<A> = A => void;
@@ -65,6 +69,27 @@ export function useContext<T>(Context: ReactContext<T>): T {
   return dispatcher.useContext(Context);
 }
 
+export function unstable_useContextWithBailout<T>(
+  context: ReactContext<T>,
+  select: (T => Array<mixed>) | null,
+): T {
+  if (!(enableLazyContextPropagation && enableContextProfiling)) {
+    throw new Error('Not implemented.');
+  }
+
+  const dispatcher = resolveDispatcher();
+  if (__DEV__) {
+    if (context.$$typeof === REACT_CONSUMER_TYPE) {
+      console.error(
+        'Calling useContext(Context.Consumer) is not supported and will cause bugs. ' +
+          'Did you mean to call useContext(Context) instead?',
+      );
+    }
+  }
+  // $FlowFixMe[not-a-function] This is unstable, thus optional
+  return dispatcher.unstable_useContextWithBailout(context, select);
+}
+
 export function useState<S>(
   initialState: (() => S) | S,
 ): [S, Dispatch<BasicStateAction<S>>] {

packages/shared/ReactFeatureFlags.js

@@ -97,6 +97,9 @@ export const enableTransitionTracing = false;
 // No known bugs, but needs performance testing
 export const enableLazyContextPropagation = false;
 
+// Expose unstable useContext for performance testing
+export const enableContextProfiling = false;
+
 // FB-only usage. The new API has different semantics.
 export const enableLegacyHidden = false;
 

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

@@ -57,6 +57,7 @@ export const enableFlightReadableStream = true;
 export const enableGetInspectorDataForInstanceInProduction = true;
 export const enableInfiniteRenderLoopDetection = true;
 export const enableLazyContextPropagation = false;
+export const enableContextProfiling = false;
 export const enableLegacyCache = false;
 export const enableLegacyFBSupport = false;
 export const enableLegacyHidden = false;

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

@@ -50,6 +50,7 @@ export const enableFlightReadableStream = true;
 export const enableGetInspectorDataForInstanceInProduction = false;
 export const enableInfiniteRenderLoopDetection = true;
 export const enableLazyContextPropagation = false;
+export const enableContextProfiling = false;
 export const enableLegacyCache = false;
 export const enableLegacyFBSupport = false;
 export const enableLegacyHidden = false;

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

@@ -52,6 +52,7 @@ export const transitionLaneExpirationMs = 5000;
 
 export const disableSchedulerTimeoutInWorkLoop = false;
 export const enableLazyContextPropagation = false;
+export const enableContextProfiling = false;
 export const enableLegacyHidden = false;
 
 export const consoleManagedByDevToolsDuringStrictMode = false;

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

@@ -42,6 +42,7 @@ export const enableFlightReadableStream = true;
 export const enableGetInspectorDataForInstanceInProduction = false;
 export const enableInfiniteRenderLoopDetection = true;
 export const enableLazyContextPropagation = false;
+export const enableContextProfiling = false;
 export const enableLegacyCache = false;
 export const enableLegacyFBSupport = false;
 export const enableLegacyHidden = false;

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

@@ -55,6 +55,7 @@ export const transitionLaneExpirationMs = 5000;
 
 export const disableSchedulerTimeoutInWorkLoop = false;
 export const enableLazyContextPropagation = false;
+export const enableContextProfiling = false;
 export const enableLegacyHidden = false;
 
 export const consoleManagedByDevToolsDuringStrictMode = false;

packages/shared/forks/ReactFeatureFlags.www.js

@@ -78,6 +78,8 @@ export const enableTaint = false;
 
 export const enablePostpone = false;
 
+export const enableContextProfiling = true;
+
 // TODO: www currently relies on this feature. It's disabled in open source.
 // Need to remove it.
 export const disableCommentsAsDOMContainers = false;

scripts/error-codes/codes.json

@@ -525,5 +525,6 @@
   "537": "Cannot pass event handlers (%s) in renderToMarkup because the HTML will never be hydrated so they can never get called.",
   "538": "Cannot use state or effect Hooks in renderToMarkup because this component will never be hydrated.",
   "539": "Binary RSC chunks cannot be encoded as strings. This is a bug in the wiring of the React streams.",
-  "540": "String chunks need to be passed in their original shape. Not split into smaller string chunks. This is a bug in the wiring of the React streams."
-}
+  "540": "String chunks need to be passed in their original shape. Not split into smaller string chunks. This is a bug in the wiring of the React streams.",
+  "541": "Compared context values must be arrays"
+}