Skip to content

Commit 0be9bfa

Browse files
gnoffgnoff
committed
Implementation of useContextSelector hook
1 parent 173b9bb commit 0be9bfa

File tree

4 files changed

+279
-1
lines changed

4 files changed

+279
-1
lines changed

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

Lines changed: 89 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
1717
import ReactSharedInternals from 'shared/ReactSharedInternals';
1818

1919
import {NoWork} from './ReactFiberExpirationTime';
20-
import {readContext} from './ReactFiberNewContext';
20+
import {readContext, selectFromContext} from './ReactFiberNewContext';
2121
import {
2222
Update as UpdateEffect,
2323
Passive as PassiveEffect,
@@ -64,6 +64,7 @@ export type Dispatcher = {
6464
context: ReactContext<T>,
6565
observedBits: void | number | boolean,
6666
): T,
67+
useContextSelector<T, S>(context: ReactContext<T>, selector: (T) => S): S,
6768
useRef<T>(initialValue: T): {current: T},
6869
useEffect(
6970
create: () => (() => void) | void,
@@ -602,6 +603,63 @@ function updateWorkInProgressHook(): Hook {
602603
return workInProgressHook;
603604
}
604605

606+
function makeSelect<T, S>(
607+
context: ReactContext<T>,
608+
selector: T => S,
609+
): (T, (T) => S) => [S, boolean] {
610+
// close over memoized value and selection
611+
let previousValue, previousSelection;
612+
613+
// select function will return a tuple of the selection as well as whether
614+
// the selection was a new value or not
615+
return function select(value: T) {
616+
let selection = previousSelection;
617+
let isNew = false;
618+
619+
// don't recompute if values are the same
620+
if (!is(value, previousValue)) {
621+
selection = selector(value);
622+
if (!is(selection, previousSelection)) {
623+
// if same we can still consider the selection memoized since the selected values are identical
624+
isNew = true;
625+
}
626+
}
627+
previousValue = value;
628+
previousSelection = selection;
629+
return [selection, isNew];
630+
};
631+
}
632+
633+
function mountContextSelector<T, S>(
634+
context: ReactContext<T>,
635+
selector: T => S,
636+
): S {
637+
const hook = mountWorkInProgressHook();
638+
let select = makeSelect(context, selector);
639+
let [selection] = selectFromContext(context, select);
640+
hook.memoizedState = [context, selector, select];
641+
return selection;
642+
}
643+
644+
function updateContextSelector<T, S>(
645+
context: ReactContext<T>,
646+
selector: T => S,
647+
): S {
648+
const hook = updateWorkInProgressHook();
649+
let [previousContext, previousSelector, previousSelect] = hook.memoizedState;
650+
651+
if (context !== previousContext || selector !== previousSelector) {
652+
// context and or selector have changed. we need to discard memoizedState
653+
// and recreate our select function
654+
let select = makeSelect(context, selector);
655+
let [selection] = selectFromContext(context, select);
656+
hook.memoizedState = [context, selector, select];
657+
return selection;
658+
} else {
659+
return selectFromContext(context, previousSelect)[0];
660+
}
661+
}
662+
605663
function createFunctionComponentUpdateQueue(): FunctionComponentUpdateQueue {
606664
return {
607665
lastEffect: null,
@@ -1223,6 +1281,7 @@ export const ContextOnlyDispatcher: Dispatcher = {
12231281

12241282
useCallback: throwInvalidHookError,
12251283
useContext: throwInvalidHookError,
1284+
useContextSelector: throwInvalidHookError,
12261285
useEffect: throwInvalidHookError,
12271286
useImperativeHandle: throwInvalidHookError,
12281287
useLayoutEffect: throwInvalidHookError,
@@ -1238,6 +1297,7 @@ const HooksDispatcherOnMount: Dispatcher = {
12381297

12391298
useCallback: mountCallback,
12401299
useContext: readContext,
1300+
useContextSelector: mountContextSelector,
12411301
useEffect: mountEffect,
12421302
useImperativeHandle: mountImperativeHandle,
12431303
useLayoutEffect: mountLayoutEffect,
@@ -1253,6 +1313,7 @@ const HooksDispatcherOnUpdate: Dispatcher = {
12531313

12541314
useCallback: updateCallback,
12551315
useContext: readContext,
1316+
useContextSelector: updateContextSelector,
12561317
useEffect: updateEffect,
12571318
useImperativeHandle: updateImperativeHandle,
12581319
useLayoutEffect: updateLayoutEffect,
@@ -1312,6 +1373,11 @@ if (__DEV__) {
13121373
mountHookTypesDev();
13131374
return readContext(context, observedBits);
13141375
},
1376+
useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
1377+
currentHookNameInDev = 'useContextSelector';
1378+
mountHookTypesDev();
1379+
return mountContextSelector(context, selector);
1380+
},
13151381
useEffect(
13161382
create: () => (() => void) | void,
13171383
deps: Array<mixed> | void | null,
@@ -1413,6 +1479,11 @@ if (__DEV__) {
14131479
updateHookTypesDev();
14141480
return readContext(context, observedBits);
14151481
},
1482+
useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
1483+
currentHookNameInDev = 'useContextSelector';
1484+
updateHookTypesDev();
1485+
return mountContextSelector(context, selector);
1486+
},
14161487
useEffect(
14171488
create: () => (() => void) | void,
14181489
deps: Array<mixed> | void | null,
@@ -1510,6 +1581,11 @@ if (__DEV__) {
15101581
updateHookTypesDev();
15111582
return readContext(context, observedBits);
15121583
},
1584+
useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
1585+
currentHookNameInDev = 'useContextSelector';
1586+
updateHookTypesDev();
1587+
return updateContextSelector(context, selector);
1588+
},
15131589
useEffect(
15141590
create: () => (() => void) | void,
15151591
deps: Array<mixed> | void | null,
@@ -1610,6 +1686,12 @@ if (__DEV__) {
16101686
mountHookTypesDev();
16111687
return readContext(context, observedBits);
16121688
},
1689+
useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
1690+
currentHookNameInDev = 'useContextSelector';
1691+
warnInvalidHookAccess();
1692+
mountHookTypesDev();
1693+
return mountContextSelector(context, selector);
1694+
},
16131695
useEffect(
16141696
create: () => (() => void) | void,
16151697
deps: Array<mixed> | void | null,
@@ -1718,6 +1800,12 @@ if (__DEV__) {
17181800
updateHookTypesDev();
17191801
return readContext(context, observedBits);
17201802
},
1803+
useContextSelector<T, S>(context: ReactContext<T>, selector: T => S): S {
1804+
currentHookNameInDev = 'useContextSelector';
1805+
warnInvalidHookAccess();
1806+
updateHookTypesDev();
1807+
return updateContextSelector(context, selector);
1808+
},
17211809
useEffect(
17221810
create: () => (() => void) | void,
17231811
deps: Array<mixed> | void | null,

‎packages/react-reconciler/src/__tests__/ReactNewContext-test.internal.js

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1355,6 +1355,165 @@ describe('ReactNewContext', () => {
13551355
});
13561356
});
13571357

1358+
describe('useContextSelector', () => {
1359+
it('context propagation defers checks as long as possible', () => {
1360+
const Context = React.createContext('abcdefg');
1361+
1362+
let lastSelector;
1363+
1364+
let i = 0;
1365+
1366+
let makeSelector = () => {
1367+
lastSelector = (j => v => {
1368+
Scheduler.yieldValue('selector' + j);
1369+
return v;
1370+
})(i++);
1371+
return lastSelector;
1372+
};
1373+
1374+
makeSelector();
1375+
1376+
let Foo = React.memo(function Foo({selector}) {
1377+
Scheduler.yieldValue('Foo');
1378+
let selection = React.useContextSelector(Context, selector);
1379+
return <span>{selection}</span>;
1380+
});
1381+
1382+
let App = ({value, selector}) => {
1383+
return (
1384+
<Context.Provider value={value}>
1385+
<div>
1386+
<Foo selector={selector} />
1387+
</div>
1388+
</Context.Provider>
1389+
);
1390+
};
1391+
1392+
// initial render
1393+
ReactNoop.render(<App value="abcdefg" selector={makeSelector()} />);
1394+
expect(Scheduler).toFlushAndYield(['Foo', 'selector1']);
1395+
1396+
// different selector -> Foo should do memo check and take new selector and then update
1397+
ReactNoop.render(<App value="abcdefgh" selector={makeSelector()} />);
1398+
expect(Scheduler).toFlushAndYield(['Foo', 'selector2']);
1399+
1400+
// shallow equal props -> memo should bailout, no selector was called but memoized so no yield
1401+
ReactNoop.render(<App value="abcdefgh" selector={lastSelector} />);
1402+
expect(Scheduler).toFlushAndYield([]);
1403+
1404+
// differe context value, memo props shallow equal
1405+
// -> call selector before attempted bailout, end up updating instead of bailout
1406+
ReactNoop.render(<App value="abcdefghi" selector={lastSelector} />);
1407+
expect(Scheduler).toFlushAndYield(['selector2', 'Foo']);
1408+
});
1409+
it('general test', () => {
1410+
const Context = React.createContext('abcdefg');
1411+
const FooContext = React.createContext(0);
1412+
const BarContext = React.createContext(0);
1413+
1414+
function Provider(props) {
1415+
return (
1416+
<Context.Provider value={props.string}>
1417+
{props.children}
1418+
</Context.Provider>
1419+
);
1420+
}
1421+
1422+
function Foo(props) {
1423+
let index = React.useContext(FooContext);
1424+
let selector = React.useCallback(v => v.substring(0, index), [index]);
1425+
let selection = React.useContextSelector(Context, selector);
1426+
Scheduler.yieldValue('Foo');
1427+
return <span prop={'foo selection: ' + selection} />;
1428+
}
1429+
1430+
function Bar(props) {
1431+
let index = React.useContext(BarContext);
1432+
let selector = React.useCallback(v => v.substring(index), [index]);
1433+
let selection = React.useContextSelector(Context, selector);
1434+
Scheduler.yieldValue('Bar');
1435+
return <span prop={'bar selection: ' + selection} />;
1436+
}
1437+
1438+
class Indirection extends React.Component {
1439+
shouldComponentUpdate() {
1440+
return false;
1441+
}
1442+
render() {
1443+
return this.props.children;
1444+
}
1445+
}
1446+
1447+
function App(props) {
1448+
return (
1449+
<FooContext.Provider value={props.fooIndex}>
1450+
<BarContext.Provider value={props.barIndex}>
1451+
<Provider string={props.string}>
1452+
<Indirection {...props}>
1453+
<Indirection>
1454+
<Foo />
1455+
</Indirection>
1456+
<Indirection>
1457+
<Bar />
1458+
</Indirection>
1459+
</Indirection>
1460+
</Provider>
1461+
</BarContext.Provider>
1462+
</FooContext.Provider>
1463+
);
1464+
}
1465+
1466+
ReactNoop.render(<App string="abcdefg" fooIndex={2} barIndex={2} />);
1467+
expect(Scheduler).toFlushAndYield(['Foo', 'Bar']);
1468+
expect(ReactNoop.getChildren()).toEqual([
1469+
span('foo selection: ab'),
1470+
span('bar selection: cdefg'),
1471+
]);
1472+
1473+
ReactNoop.render(<App string="abcdefg" fooIndex={3} barIndex={2} />);
1474+
expect(Scheduler).toFlushAndYield(['Foo']);
1475+
expect(ReactNoop.getChildren()).toEqual([
1476+
span('foo selection: abc'),
1477+
span('bar selection: cdefg'),
1478+
]);
1479+
1480+
ReactNoop.render(<App string="a*cdefg" fooIndex={3} barIndex={2} />);
1481+
expect(Scheduler).toFlushAndYield(['Foo']);
1482+
expect(ReactNoop.getChildren()).toEqual([
1483+
span('foo selection: a*c'),
1484+
span('bar selection: cdefg'),
1485+
]);
1486+
1487+
ReactNoop.render(<App string="a*cdefg" fooIndex={3} barIndex={1} />);
1488+
expect(Scheduler).toFlushAndYield(['Bar']);
1489+
expect(ReactNoop.getChildren()).toEqual([
1490+
span('foo selection: a*c'),
1491+
span('bar selection: *cdefg'),
1492+
]);
1493+
1494+
ReactNoop.render(<App string="a|cdefg" fooIndex={3} barIndex={1} />);
1495+
expect(Scheduler).toFlushAndYield(['Foo', 'Bar']);
1496+
expect(ReactNoop.getChildren()).toEqual([
1497+
span('foo selection: a|c'),
1498+
span('bar selection: |cdefg'),
1499+
]);
1500+
1501+
ReactNoop.render(<App string="a|cdefg" fooIndex={3} barIndex={4} />);
1502+
expect(Scheduler).toFlushAndYield(['Bar']);
1503+
expect(ReactNoop.getChildren()).toEqual([
1504+
span('foo selection: a|c'),
1505+
span('bar selection: efg'),
1506+
]);
1507+
1508+
ReactNoop.render(<App string="a|c-efg" fooIndex={3} barIndex={4} />);
1509+
expect(Scheduler).toFlushAndYield([]);
1510+
expect(ReactNoop.getChildren()).toEqual([
1511+
span('foo selection: a|c'),
1512+
span('bar selection: efg'),
1513+
]);
1514+
});
1515+
});
1516+
13581517
describe('Context.Consumer', () => {
13591518
it('warns if child is not a function', () => {
13601519
spyOnDev(console, 'error');

‎packages/react/src/React.js

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import memo from './memo';
3131
import {
3232
useCallback,
3333
useContext,
34+
useContextSelector,
3435
useEffect,
3536
useImperativeHandle,
3637
useDebugValue,
@@ -75,6 +76,7 @@ const React = {
7576

7677
useCallback,
7778
useContext,
79+
useContextSelector,
7880
useEffect,
7981
useImperativeHandle,
8082
useDebugValue,

‎packages/react/src/ReactHooks.js

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,35 @@ export function useContext<T>(
6969
return dispatcher.useContext(Context, unstable_observedBits);
7070
}
7171

72+
export function useContextSelector<T, S>(
73+
Context: ReactContext<T>,
74+
selector: T => S,
75+
) {
76+
const dispatcher = resolveDispatcher();
77+
if (__DEV__) {
78+
// TODO: add a more generic warning for invalid values.
79+
if ((Context: any)._context !== undefined) {
80+
const realContext = (Context: any)._context;
81+
// Don't deduplicate because this legitimately causes bugs
82+
// and nobody should be using this in existing code.
83+
if (realContext.Consumer === Context) {
84+
warning(
85+
false,
86+
'Calling useContextSelector(Context.Consumer, selector) is not supported, may cause bugs, and will be ' +
87+
'removed in a future major release. Did you mean to call useContextSelector(Context, selector) instead?',
88+
);
89+
} else if (realContext.Provider === Context) {
90+
warning(
91+
false,
92+
'Calling useContext(Context.Provider, selector) is not supported. ' +
93+
'Did you mean to call useContext(Contextm, selector) instead?',
94+
);
95+
}
96+
}
97+
}
98+
return dispatcher.useContextSelector(Context, selector);
99+
}
100+
72101
export function useState<S>(initialState: (() => S) | S) {
73102
const dispatcher = resolveDispatcher();
74103
return dispatcher.useState(initialState);

0 commit comments

Comments
 (0)