Implementation of useContextSelector hook

Author: gnoff

packages/react-reconciler/src/ReactFiberHooks.js

@@ -17,7 +17,7 @@ import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 
 import {NoWork} from './ReactFiberExpirationTime';
-import {readContext} from './ReactFiberNewContext';
+import {readContext, selectFromContext} from './ReactFiberNewContext';
 import {
   Update as UpdateEffect,
   Passive as PassiveEffect,
@@ -64,6 +64,7 @@ export type Dispatcher = {
     context: ReactContext<T>,
     observedBits: void | number | boolean,
   ): T,
+  useContextSelector<T, S>(context: ReactContext<T>, selector: (T) => S): S,
   useRef<T>(initialValue: T): {current: T},
   useEffect(
     create: () => (() => void) | void,
@@ -602,6 +603,63 @@ function updateWorkInProgressHook(): Hook {
   return workInProgressHook;
 }
 
+function makeSelect<T, S>(
+  context: ReactContext<T>,
+  selector: T => S,
+): (T, (T) => S) => [S, boolean] {
+  // close over memoized value and selection
+  let previousValue, previousSelection;
+
+  // select function will return a tuple of the selection as well as whether
+  // the selection was a new value or not
+  return function select(value: T) {
+    let selection = previousSelection;
+    let isNew = false;
+
+    // don't recompute if values are the same
+    if (!is(value, previousValue)) {
+      selection = selector(value);
+      if (!is(selection, previousSelection)) {
+        // if same we can still consider the selection memoized since the selected values are identical
+        isNew = true;
+      }
+    }
+    previousValue = value;
+    previousSelection = selection;
+    return [selection, isNew];
+  };
+}
+
+function mountContextSelector<T, S>(
+  context: ReactContext<T>,
+  selector: T => S,
+): S {
+  const hook = mountWorkInProgressHook();
+  let select = makeSelect(context, selector);
+  let [selection] = selectFromContext(context, select);
+  hook.memoizedState = [context, selector, select];
+  return selection;
+}
+
+function updateContextSelector<T, S>(
+  context: ReactContext<T>,
+  selector: T => S,
+): S {
+  const hook = updateWorkInProgressHook();
+  let [previousContext, previousSelector, previousSelect] = hook.memoizedState;
+
+  if (context !== previousContext || selector !== previousSelector) {
+    // context and or selector have changed. we need to discard memoizedState
+    // and recreate our select function
+    let select = makeSelect(context, selector);
+    let [selection] = selectFromContext(context, select);
+    hook.memoizedState = [context, selector, select];
+    return selection;
+  } else {
+    return selectFromContext(context, previousSelect)[0];
+  }
+}
+
 function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
   return {
     lastEffect: null,
@@ -1223,6 +1281,7 @@ export const ContextOnlyDispatcher: Dispatcher = {
 
   useCallback: throwInvalidHookError,
   useContext: throwInvalidHookError,
+  useContextSelector: throwInvalidHookError,
   useEffect: throwInvalidHookError,
   useImperativeHandle: throwInvalidHookError,
   useLayoutEffect: throwInvalidHookError,
@@ -1238,6 +1297,7 @@ const HooksDispatcherOnMount: Dispatcher = {
 
   useCallback: mountCallback,
   useContext: readContext,
+  useContextSelector: mountContextSelector,
   useEffect: mountEffect,
   useImperativeHandle: mountImperativeHandle,
   useLayoutEffect: mountLayoutEffect,
@@ -1253,6 +1313,7 @@ const HooksDispatcherOnUpdate: Dispatcher = {
 
   useCallback: updateCallback,
   useContext: readContext,
+  useContextSelector: updateContextSelector,
   useEffect: updateEffect,
   useImperativeHandle: updateImperativeHandle,
   useLayoutEffect: updateLayoutEffect,
@@ -1312,6 +1373,11 @@ if (__DEV__) {
       mountHookTypesDev();
       return readContext(context, observedBits);
     },
+    useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
+      currentHookNameInDev = 'useContextSelector';
+      mountHookTypesDev();
+      return mountContextSelector(context, selector);
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -1413,6 +1479,11 @@ if (__DEV__) {
       updateHookTypesDev();
       return readContext(context, observedBits);
     },
+    useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
+      currentHookNameInDev = 'useContextSelector';
+      updateHookTypesDev();
+      return mountContextSelector(context, selector);
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -1510,6 +1581,11 @@ if (__DEV__) {
       updateHookTypesDev();
       return readContext(context, observedBits);
     },
+    useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
+      currentHookNameInDev = 'useContextSelector';
+      updateHookTypesDev();
+      return updateContextSelector(context, selector);
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -1610,6 +1686,12 @@ if (__DEV__) {
       mountHookTypesDev();
       return readContext(context, observedBits);
     },
+    useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
+      currentHookNameInDev = 'useContextSelector';
+      warnInvalidHookAccess();
+      mountHookTypesDev();
+      return mountContextSelector(context, selector);
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,
@@ -1718,6 +1800,12 @@ if (__DEV__) {
       updateHookTypesDev();
       return readContext(context, observedBits);
     },
+    useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
+      currentHookNameInDev = 'useContextSelector';
+      warnInvalidHookAccess();
+      updateHookTypesDev();
+      return updateContextSelector(context, selector);
+    },
     useEffect(
       create: () => (() => void) | void,
       deps: Array<mixed> | void | null,

packages/react-reconciler/src/ReactFiberNewContext.js

@@ -266,7 +266,21 @@ export function checkContextDependencies(
         let context = dependency.context;
         let observedBits = dependency.observedBits;
         if ((observedBits & context._currentChangedBits) !== 0) {
-          return true;
+          let requiresUpdate = true;
+
+          let selector = dependency.selector;
+          if (typeof selector === 'function') {
+            let [, isNew] = selector(
+              isPrimaryRenderer
+                ? context._currentValue
+                : context._currentValue2,
+            );
+            requiresUpdate = isNew;
+          }
+
+          if (requiresUpdate) {
+            return true;
+          }
         }
         dependency = dependency.next;
       }
@@ -634,3 +648,50 @@ export function readContext<T>(
   }
   return isPrimaryRenderer ? context._currentValue : context._currentValue2;
 }
+
+export function selectFromContext<T, S>(
+  context: ReactContext<T>,
+  select: T => [S, boolean],
+): [S, boolean] {
+  if (__DEV__) {
+    // This warning would fire if you read context inside a Hook like useMemo.
+    // Unlike the class check below, it's not enforced in production for perf.
+    warning(
+      !isDisallowedContextReadInDEV,
+      '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().',
+    );
+  }
+
+  let contextItem = {
+    context: ((context: any): ReactContext<mixed>),
+    observedBits: MAX_SIGNED_31_BIT_INT,
+    selector: select,
+    next: null,
+  };
+
+  if (lastContextDependency === null) {
+    invariant(
+      currentlyRenderingFiber !== null,
+      '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;
+    currentlyRenderingFiber.contextDependencies = {
+      first: contextItem,
+      expirationTime: NoWork,
+    };
+  } else {
+    // Append a new context item.
+    lastContextDependency = lastContextDependency.next = contextItem;
+  }
+  return isPrimaryRenderer
+    ? select(context._currentValue)
+    : select(context._currentValue2);
+}