Implement naive version of context selectors

For internal experimentation only.

This implements `unstable_useSelectedContext` behind a feature flag.
It's based on [RFC 119](reactjs/rfcs#119) and
[RFC 118](reactjs/rfcs#118) by @gnoff.

Usage:

```js
const selection = useSelectedContext(Context, c => select(c));
```

The key feature is that if the selected value does not change between
renders, the component will bail out of rendering its children, a la
`memo`, `PureComponent`, or the `useState` bailout mechanism. (Unless
some other state, props, or context was updated in the same render.)

However, I have not implemented the RFC's proposed optimizations to
context propagation. We would like to land those eventually, but doing
so will require a refactor that we don't currently have the bandwidth to
complete. It will need to wait until after React 18.

In the meantime though, we believe there may be value in landing this
more naive implementation. It's designed to be API-compatible with the
full proposal, so we have the option to make those optimizations in
a non-breaking release. However, since it's still behind a flag, this
currently has no impact on the stable release channel. We reserve the
right to change or remove it, as we conduct internal testing.

I also added an optional third argument, `isSelectionEqual`. If defined,
it will override the default comparison function used to check if the
selected value has changed (`Object.is`).

Author: acdlite

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

@@ -129,6 +129,21 @@ function useContext<T>(
   return context._currentValue;
 }
 
+function useSelectedContext<C, S>(
+  Context: ReactContext<C>,
+  selector: C => S,
+  isEqual: ((S, S) => boolean) | void,
+): S {
+  const context = Context._currentValue;
+  const selection = selector(context);
+  hookLog.push({
+    primitive: 'SelectedContext',
+    stackError: new Error(),
+    value: selection,
+  });
+  return selection;
+}
+
 function useState<S>(
   initialState: (() => S) | S,
 ): [S, Dispatch<BasicStateAction<S>>] {
@@ -322,6 +337,7 @@ const Dispatcher: DispatcherType = {
   useCacheRefresh,
   useCallback,
   useContext,
+  useSelectedContext,
   useEffect,
   useImperativeHandle,
   useDebugValue,

packages/react-dom/src/server/ReactPartialRendererHooks.js

@@ -251,6 +251,22 @@ function useContext<T>(
   return context[threadID];
 }
 
+function useSelectedContext<C, S>(
+  Context: ReactContext<C>,
+  selector: C => S,
+  isEqual: ((S, S) => boolean) | void,
+): S {
+  if (__DEV__) {
+    currentHookNameInDev = 'useSelectedContext';
+  }
+  resolveCurrentlyRenderingComponent();
+  const threadID = currentPartialRenderer.threadID;
+  validateContextBounds(Context, threadID);
+  const context = Context[threadID];
+  const selection = selector(context);
+  return selection;
+}
+
 function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {
   // $FlowFixMe: Flow doesn't like mixed types
   return typeof action === 'function' ? action(state) : action;
@@ -503,6 +519,7 @@ export function setCurrentPartialRenderer(renderer: PartialRenderer) {
 export const Dispatcher: DispatcherType = {
   readContext,
   useContext,
+  useSelectedContext,
   useMemo,
   useReducer,
   useRef,

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

@@ -30,6 +30,7 @@ import {
   decoupleUpdatePriorityFromScheduler,
   enableUseRefAccessWarning,
   enableDoubleInvokingEffects,
+  enableContextSelectors,
 } from 'shared/ReactFeatureFlags';
 
 import {
@@ -52,7 +53,7 @@ import {
   higherLanePriority,
   DefaultLanePriority,
 } from './ReactFiberLane.new';
-import {readContext} from './ReactFiberNewContext.new';
+import {readContext, readContextInsideHook} from './ReactFiberNewContext.new';
 import {HostRoot, CacheComponent} from './ReactWorkTags';
 import {
   Update as UpdateEffect,
@@ -627,6 +628,56 @@ function updateWorkInProgressHook(): Hook {
   return workInProgressHook;
 }
 
+function mountSelectedContext<C, S>(
+  Context: ReactContext<C>,
+  selector: C => S,
+  isEqual: ((S, S) => boolean) | void,
+): S {
+  if (!enableContextSelectors) {
+    return (undefined: any);
+  }
+
+  const hook = mountWorkInProgressHook();
+  const context = readContextInsideHook(Context);
+  const selection = selector(context);
+  hook.memoizedState = selection;
+  return selection;
+}
+
+function updateSelectedContext<C, S>(
+  Context: ReactContext<C>,
+  selector: C => S,
+  isEqual: ((S, S) => boolean) | void,
+): S {
+  if (!enableContextSelectors) {
+    return (undefined: any);
+  }
+
+  const hook = updateWorkInProgressHook();
+  const context = readContextInsideHook(Context);
+  const newSelection = selector(context);
+  const oldSelection: S = hook.memoizedState;
+  if (isEqual !== undefined) {
+    if (__DEV__) {
+      if (typeof isEqual !== 'function') {
+        console.error(
+          'The optional third argument to useSelectedContext must be a ' +
+            'function. Instead got: %s',
+          isEqual,
+        );
+      }
+    }
+    if (isEqual(newSelection, oldSelection)) {
+      return oldSelection;
+    }
+  } else if (is(newSelection, oldSelection)) {
+    return oldSelection;
+  }
+  markWorkInProgressReceivedUpdate();
+  hook.memoizedState = newSelection;
+  return newSelection;
+}
+
 function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
   return {
     lastEffect: null,
@@ -1995,6 +2046,7 @@ export const ContextOnlyDispatcher: Dispatcher = {
 
   useCallback: throwInvalidHookError,
   useContext: throwInvalidHookError,
+  useSelectedContext: throwInvalidHookError,
   useEffect: throwInvalidHookError,
   useImperativeHandle: throwInvalidHookError,
   useLayoutEffect: throwInvalidHookError,
@@ -2020,6 +2072,7 @@ const HooksDispatcherOnMount: Dispatcher = {
 
   useCallback: mountCallback,
   useContext: readContext,
+  useSelectedContext: mountSelectedContext,
   useEffect: mountEffect,
   useImperativeHandle: mountImperativeHandle,
   useLayoutEffect: mountLayoutEffect,
@@ -2045,6 +2098,7 @@ const HooksDispatcherOnUpdate: Dispatcher = {
 
   useCallback: updateCallback,
   useContext: readContext,
+  useSelectedContext: updateSelectedContext,
   useEffect: updateEffect,
   useImperativeHandle: updateImperativeHandle,
   useLayoutEffect: updateLayoutEffect,
@@ -2070,6 +2124,7 @@ const HooksDispatcherOnRerender: Dispatcher = {
 
   useCallback: updateCallback,
   useContext: readContext,
+  useSelectedContext: updateSelectedContext,
   useEffect: updateEffect,
   useImperativeHandle: updateImperativeHandle,
   useLayoutEffect: updateLayoutEffect,
@@ -2138,6 +2193,21 @@ if (__DEV__) {
       mountHookTypesDev();
       return readContext(context, observedBits);
     },
+    useSelectedContext<C, S>(
+      context: ReactContext<C>,
+      selector: C => S,
+      isEqual: ((S, S) => boolean) | void,
+    ): S {
+      currentHookNameInDev = 'useSelectedContext';
+      mountHookTypesDev();
+      const prevDispatcher = ReactCurrentDispatcher.current;
+      ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnMountInDEV;
+      try {
+        return mountSelectedContext(context, selector, isEqual);
+      } finally {
+        ReactCurrentDispatcher.current = prevDispatcher;
+      }
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -2272,6 +2342,21 @@ if (__DEV__) {
       updateHookTypesDev();
       return readContext(context, observedBits);
     },
+    useSelectedContext<C, S>(
+      context: ReactContext<C>,
+      selector: C => S,
+      isEqual: ((S, S) => boolean) | void,
+    ): S {
+      currentHookNameInDev = 'useSelectedContext';
+      updateHookTypesDev();
+      const prevDispatcher = ReactCurrentDispatcher.current;
+      ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnMountInDEV;
+      try {
+        return mountSelectedContext(context, selector, isEqual);
+      } finally {
+        ReactCurrentDispatcher.current = prevDispatcher;
+      }
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -2402,6 +2487,21 @@ if (__DEV__) {
       updateHookTypesDev();
       return readContext(context, observedBits);
     },
+    useSelectedContext<C, S>(
+      context: ReactContext<C>,
+      selector: C => S,
+      isEqual: ((S, S) => boolean) | void,
+    ): S {
+      currentHookNameInDev = 'useSelectedContext';
+      updateHookTypesDev();
+      const prevDispatcher = ReactCurrentDispatcher.current;
+      ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnUpdateInDEV;
+      try {
+        return updateSelectedContext(context, selector, isEqual);
+      } finally {
+        ReactCurrentDispatcher.current = prevDispatcher;
+      }
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -2533,6 +2633,21 @@ if (__DEV__) {
       updateHookTypesDev();
       return readContext(context, observedBits);
     },
+    useSelectedContext<C, S>(
+      context: ReactContext<C>,
+      selector: C => S,
+      isEqual: ((S, S) => boolean) | void,
+    ): S {
+      currentHookNameInDev = 'useSelectedContext';
+      updateHookTypesDev();
+      const prevDispatcher = ReactCurrentDispatcher.current;
+      ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnRerenderInDEV;
+      try {
+        return updateSelectedContext(context, selector, isEqual);
+      } finally {
+        ReactCurrentDispatcher.current = prevDispatcher;
+      }
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -2666,6 +2781,22 @@ if (__DEV__) {
       mountHookTypesDev();
       return readContext(context, observedBits);
     },
+    useSelectedContext<C, S>(
+      context: ReactContext<C>,
+      selector: C => S,
+      isEqual: ((S, S) => boolean) | void,
+    ): S {
+      currentHookNameInDev = 'useSelectedContext';
+      warnInvalidHookAccess();
+      mountHookTypesDev();
+      const prevDispatcher = ReactCurrentDispatcher.current;
+      ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnMountInDEV;
+      try {
+        return mountSelectedContext(context, selector, isEqual);
+      } finally {
+        ReactCurrentDispatcher.current = prevDispatcher;
+      }
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -2811,6 +2942,22 @@ if (__DEV__) {
       updateHookTypesDev();
       return readContext(context, observedBits);
     },
+    useSelectedContext<C, S>(
+      context: ReactContext<C>,
+      selector: C => S,
+      isEqual: ((S, S) => boolean) | void,
+    ): S {
+      currentHookNameInDev = 'useSelectedContext';
+      warnInvalidHookAccess();
+      updateHookTypesDev();
+      const prevDispatcher = ReactCurrentDispatcher.current;
+      ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnUpdateInDEV;
+      try {
+        return updateSelectedContext(context, selector, isEqual);
+      } finally {
+        ReactCurrentDispatcher.current = prevDispatcher;
+      }
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -2957,6 +3104,22 @@ if (__DEV__) {
       updateHookTypesDev();
       return readContext(context, observedBits);
     },
+    useSelectedContext<C, S>(
+      context: ReactContext<C>,
+      selector: C => S,
+      isEqual: ((S, S) => boolean) | void,
+    ): S {
+      currentHookNameInDev = 'useSelectedContext';
+      warnInvalidHookAccess();
+      updateHookTypesDev();
+      const prevDispatcher = ReactCurrentDispatcher.current;
+      ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnUpdateInDEV;
+      try {
+        return updateSelectedContext(context, selector, isEqual);
+      } finally {
+        ReactCurrentDispatcher.current = prevDispatcher;
+      }
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,