Implementation of useContextSelector hook

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/__tests__/ReactNewContext-test.internal.js

@@ -1355,6 +1355,165 @@ describe('ReactNewContext', () => {
     });
   });
 
+  describe('useContextSelector', () => {
+    it('context propagation defers checks as long as possible', () => {
+      const Context = React.createContext('abcdefg');
+
+      let lastSelector;
+
+      let i = 0;
+
+      let makeSelector = () => {
+        lastSelector = (j => v => {
+          Scheduler.yieldValue('selector' + j);
+          return v;
+        })(i++);
+        return lastSelector;
+      };
+
+      makeSelector();
+
+      let Foo = React.memo(function Foo({selector}) {
+        Scheduler.yieldValue('Foo');
+        let selection = React.useContextSelector(Context, selector);
+        return <span>{selection}</span>;
+      });
+
+      let App = ({value, selector}) => {
+        return (
+          <Context.Provider value={value}>
+            <div>
+              <Foo selector={selector} />
+            </div>
+          </Context.Provider>
+        );
+      };
+
+      // initial render
+      ReactNoop.render(<App value="abcdefg" selector={makeSelector()} />);
+      expect(Scheduler).toFlushAndYield(['Foo', 'selector1']);
+
+      // different selector -> Foo should do memo check and take new selector and then update
+      ReactNoop.render(<App value="abcdefgh" selector={makeSelector()} />);
+      expect(Scheduler).toFlushAndYield(['Foo', 'selector2']);
+
+      // shallow equal props -> memo should bailout, no selector was called but memoized so no yield
+      ReactNoop.render(<App value="abcdefgh" selector={lastSelector} />);
+      expect(Scheduler).toFlushAndYield([]);
+
+      // differe context value, memo props shallow equal
+      // -> call selector before attempted bailout, end up updating instead of bailout
+      ReactNoop.render(<App value="abcdefghi" selector={lastSelector} />);
+      expect(Scheduler).toFlushAndYield(['selector2', 'Foo']);
+    });
+    it('general test', () => {
+      const Context = React.createContext('abcdefg');
+      const FooContext = React.createContext(0);
+      const BarContext = React.createContext(0);
+
+      function Provider(props) {
+        return (
+          <Context.Provider value={props.string}>
+            {props.children}
+          </Context.Provider>
+        );
+      }
+
+      function Foo(props) {
+        let index = React.useContext(FooContext);
+        let selector = React.useCallback(v => v.substring(0, index), [index]);
+        let selection = React.useContextSelector(Context, selector);
+        Scheduler.yieldValue('Foo');
+        return <span prop={'foo selection: ' + selection} />;
+      }
+
+      function Bar(props) {
+        let index = React.useContext(BarContext);
+        let selector = React.useCallback(v => v.substring(index), [index]);
+        let selection = React.useContextSelector(Context, selector);
+        Scheduler.yieldValue('Bar');
+        return <span prop={'bar selection: ' + selection} />;
+      }
+
+      class Indirection extends React.Component {
+        shouldComponentUpdate() {
+          return false;
+        }
+        render() {
+          return this.props.children;
+        }
+      }
+
+      function App(props) {
+        return (
+          <FooContext.Provider value={props.fooIndex}>
+            <BarContext.Provider value={props.barIndex}>
+              <Provider string={props.string}>
+                <Indirection {...props}>
+                  <Indirection>
+                    <Foo />
+                  </Indirection>
+                  <Indirection>
+                    <Bar />
+                  </Indirection>
+                </Indirection>
+              </Provider>
+            </BarContext.Provider>
+          </FooContext.Provider>
+        );
+      }
+
+      ReactNoop.render(<App string="abcdefg" fooIndex={2} barIndex={2} />);
+      expect(Scheduler).toFlushAndYield(['Foo', 'Bar']);
+      expect(ReactNoop.getChildren()).toEqual([
+        span('foo selection: ab'),
+        span('bar selection: cdefg'),
+      ]);
+
+      ReactNoop.render(<App string="abcdefg" fooIndex={3} barIndex={2} />);
+      expect(Scheduler).toFlushAndYield(['Foo']);
+      expect(ReactNoop.getChildren()).toEqual([
+        span('foo selection: abc'),
+        span('bar selection: cdefg'),
+      ]);
+
+      ReactNoop.render(<App string="a*cdefg" fooIndex={3} barIndex={2} />);
+      expect(Scheduler).toFlushAndYield(['Foo']);
+      expect(ReactNoop.getChildren()).toEqual([
+        span('foo selection: a*c'),
+        span('bar selection: cdefg'),
+      ]);
+
+      ReactNoop.render(<App string="a*cdefg" fooIndex={3} barIndex={1} />);
+      expect(Scheduler).toFlushAndYield(['Bar']);
+      expect(ReactNoop.getChildren()).toEqual([
+        span('foo selection: a*c'),
+        span('bar selection: *cdefg'),
+      ]);
+
+      ReactNoop.render(<App string="a|cdefg" fooIndex={3} barIndex={1} />);
+      expect(Scheduler).toFlushAndYield(['Foo', 'Bar']);
+      expect(ReactNoop.getChildren()).toEqual([
+        span('foo selection: a|c'),
+        span('bar selection: |cdefg'),
+      ]);
+
+      ReactNoop.render(<App string="a|cdefg" fooIndex={3} barIndex={4} />);
+      expect(Scheduler).toFlushAndYield(['Bar']);
+      expect(ReactNoop.getChildren()).toEqual([
+        span('foo selection: a|c'),
+        span('bar selection: efg'),
+      ]);
+
+      ReactNoop.render(<App string="a|c-efg" fooIndex={3} barIndex={4} />);
+      expect(Scheduler).toFlushAndYield([]);
+      expect(ReactNoop.getChildren()).toEqual([
+        span('foo selection: a|c'),
+        span('bar selection: efg'),
+      ]);
+    });
+  });
+
   describe('Context.Consumer', () => {
     it('warns if child is not a function', () => {
       spyOnDev(console, 'error');

packages/react/src/React.js

@@ -31,6 +31,7 @@ import memo from './memo';
 import {
   useCallback,
   useContext,
+  useContextSelector,
   useEffect,
   useImperativeHandle,
   useDebugValue,
@@ -75,6 +76,7 @@ const React = {
 
   useCallback,
   useContext,
+  useContextSelector,
   useEffect,
   useImperativeHandle,
   useDebugValue,

packages/react/src/ReactHooks.js

@@ -69,6 +69,35 @@ export function useContext<T>(
   return dispatcher.useContext(Context, unstable_observedBits);
 }
 
+export function useContextSelector<T, S>(
+  Context: ReactContext<T>,
+  selector: T => S,
+) {
+  const dispatcher = resolveDispatcher();
+  if (__DEV__) {
+    // TODO: add a more generic warning for invalid values.
+    if ((Context: any)._context !== undefined) {
+      const realContext = (Context: any)._context;
+      // Don't deduplicate because this legitimately causes bugs
+      // and nobody should be using this in existing code.
+      if (realContext.Consumer === Context) {
+        warning(
+          false,
+          'Calling useContextSelector(Context.Consumer, selector) is not supported, may cause bugs, and will be ' +
+            'removed in a future major release. Did you mean to call useContextSelector(Context, selector) instead?',
+        );
+      } else if (realContext.Provider === Context) {
+        warning(
+          false,
+          'Calling useContext(Context.Provider, selector) is not supported. ' +
+            'Did you mean to call useContext(Contextm, selector) instead?',
+        );
+      }
+    }
+  }
+  return dispatcher.useContextSelector(Context, selector);
+}
+
 export function useState<S>(initialState: (() => S) | S) {
   const dispatcher = resolveDispatcher();
   return dispatcher.useState(initialState);