Skip to content

Commit

Permalink
Fork readContextForConsumer
Browse files Browse the repository at this point in the history
  • Loading branch information
Jack Pope committed Jul 23, 2024
1 parent c6a3e4d commit b3885c2
Show file tree
Hide file tree
Showing 5 changed files with 108 additions and 36 deletions.
6 changes: 5 additions & 1 deletion packages/react-debug-tools/src/ReactDebugHooks.js
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ import {
REACT_CONTEXT_TYPE,
} from 'shared/ReactSymbols';
import hasOwnProperty from 'shared/hasOwnProperty';
import type {ContextDependencyWithCompare} from '../../react-reconciler/src/ReactInternalTypes';

type CurrentDispatcherRef = typeof ReactSharedInternals;

Expand Down Expand Up @@ -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, mixed> = null;
let currentContextDependency:
| null
| ContextDependency<mixed>
| ContextDependencyWithCompare<mixed, mixed> = null;

function nextHook(): null | Hook {
const hook = currentHook;
Expand Down
16 changes: 8 additions & 8 deletions packages/react-reconciler/src/ReactFiberHooks.js
Original file line number Diff line number Diff line change
Expand Up @@ -1060,7 +1060,7 @@ function updateWorkInProgressHook(): Hook {

function unstable_useContextWithBailout<T>(
context: ReactContext<T>,
compare: void | (T => mixed),
compare: (T => mixed) | null,
): T {
return readContextAndCompare(context, compare);
}
Expand Down Expand Up @@ -4049,7 +4049,7 @@ if (__DEV__) {
}
if (enableContextProfiling) {
(HooksDispatcherOnMountInDEV: Dispatcher).unstable_useContextWithBailout =
function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
currentHookNameInDev = 'useContext';
mountHookTypesDev();
return unstable_useContextWithBailout(context, compare);
Expand Down Expand Up @@ -4238,7 +4238,7 @@ if (__DEV__) {
}
if (enableContextProfiling) {
(HooksDispatcherOnMountWithHookTypesInDEV: Dispatcher).unstable_useContextWithBailout =
function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
currentHookNameInDev = 'useContext';
updateHookTypesDev();
return unstable_useContextWithBailout(context, compare);
Expand Down Expand Up @@ -4426,7 +4426,7 @@ if (__DEV__) {
}
if (enableContextProfiling) {
(HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
currentHookNameInDev = 'useContext';
updateHookTypesDev();
return unstable_useContextWithBailout(context, compare);
Expand Down Expand Up @@ -4614,7 +4614,7 @@ if (__DEV__) {
}
if (enableContextProfiling) {
(HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
currentHookNameInDev = 'useContext';
updateHookTypesDev();
return unstable_useContextWithBailout(context, compare);
Expand Down Expand Up @@ -4828,7 +4828,7 @@ if (__DEV__) {
}
if (enableContextProfiling) {
(HooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
currentHookNameInDev = 'useContext';
warnInvalidHookAccess();
mountHookTypesDev();
Expand Down Expand Up @@ -5043,7 +5043,7 @@ if (__DEV__) {
}
if (enableContextProfiling) {
(InvalidNestedHooksDispatcherOnUpdateInDEV: Dispatcher).unstable_useContextWithBailout =
function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
currentHookNameInDev = 'useContext';
warnInvalidHookAccess();
updateHookTypesDev();
Expand Down Expand Up @@ -5258,7 +5258,7 @@ if (__DEV__) {
}
if (enableContextProfiling) {
(InvalidNestedHooksDispatcherOnRerenderInDEV: Dispatcher).unstable_useContextWithBailout =
function <T>(context: ReactContext<T>, compare: void | (T => mixed)): T {
function <T>(context: ReactContext<T>, compare: (T => mixed) | null): T {
currentHookNameInDev = 'useContext';
warnInvalidHookAccess();
updateHookTypesDev();
Expand Down
94 changes: 74 additions & 20 deletions packages/react-reconciler/src/ReactFiberNewContext.js
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import type {
Fiber,
ContextDependency,
Dependencies,
ContextDependencyWithCompare,
} from './ReactInternalTypes';
import type {StackCursor} from './ReactFiberStack';
import type {Lanes} from './ReactFiberLane';
Expand Down Expand Up @@ -72,7 +73,10 @@ if (__DEV__) {
}

let currentlyRenderingFiber: Fiber | null = null;
let lastContextDependency: ContextDependency<mixed, mixed> | null = null;
let lastContextDependency:
| ContextDependency<mixed>
| ContextDependencyWithCompare<mixed, mixed>
| null = null;
let lastFullyObservedContext: ReactContext<any> | null = null;

let isDisallowedContextReadInDEV: boolean = false;
Expand Down Expand Up @@ -403,19 +407,21 @@ function propagateContextChanges<T>(
const context: ReactContext<T> = contexts[i];
// Check if the context matches.
if (dependency.context === context) {
const compare = dependency.compare;
if (enableContextProfiling && compare != null) {
const newValue = isPrimaryRenderer
? dependency.context._currentValue
: dependency.context._currentValue2;
if (
!checkIfComparedContextValuesChanged(
dependency.lastComparedValue,
compare(newValue),
)
) {
// Compared value hasn't changed. Bail out early.
continue findContext;
if (enableContextProfiling) {
const compare = dependency.compare;
if (compare != null) {
const newValue = isPrimaryRenderer
? dependency.context._currentValue
: dependency.context._currentValue2;
if (
!checkIfComparedContextValuesChanged(
dependency.lastComparedValue,
compare(newValue),
)
) {
// Compared value hasn't changed. Bail out early.
continue findContext;
}
}
}
// Match! Schedule an update on this fiber.
Expand Down Expand Up @@ -746,13 +752,17 @@ export function prepareToReadContext(

export function readContextAndCompare<C>(
context: ReactContext<C>,
compare: void | (C => mixed),
compare: (C => mixed) | null,
): C {
if (!enableLazyContextPropagation) {
return readContext(context);
}

return readContextForConsumer(currentlyRenderingFiber, context, compare);
return readContextForConsumer_withCompare(
currentlyRenderingFiber,
context,
compare,
);
}

export function readContext<T>(context: ReactContext<T>): T {
Expand Down Expand Up @@ -782,12 +792,12 @@ export function readContextDuringReconciliation<T>(
return readContextForConsumer(consumer, context);
}

type ContextCompare<C, S> = C => S;
type ContextCompare<C, V> = C => V | null;

function readContextForConsumer<C, S>(
function readContextForConsumer_withCompare<C, S>(
consumer: Fiber | null,
context: ReactContext<C>,
compare?: void | (C => S),
compare: (C => S) | null,
): C {
const value = isPrimaryRenderer
? context._currentValue
Expand All @@ -800,7 +810,7 @@ function readContextForConsumer<C, S>(
context: ((context: any): ReactContext<mixed>),
memoizedValue: value,
next: null,
compare: ((compare: any): ContextCompare<mixed, mixed> | null),
compare: compare ? ((compare: any): ContextCompare<mixed, mixed>) : null,
lastComparedValue: compare != null ? compare(value) : null,
};

Expand Down Expand Up @@ -830,3 +840,47 @@ function readContextForConsumer<C, S>(
}
return value;
}

function readContextForConsumer<C>(
consumer: Fiber | null,
context: ReactContext<C>,
): 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,
};

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;
}
26 changes: 20 additions & 6 deletions packages/react-reconciler/src/ReactInternalTypes.js
Original file line number Diff line number Diff line change
Expand Up @@ -61,18 +61,32 @@ export type HookType =
| 'useFormState'
| 'useActionState';

export type ContextDependency<C, S> = {
export type ContextDependency<C> = {
context: ReactContext<C>,
next: ContextDependency<mixed, mixed> | null,
next:
| ContextDependency<mixed>
| ContextDependencyWithCompare<mixed, mixed>
| null,
memoizedValue: C,
};

export type ContextDependencyWithCompare<C, S> = {
context: ReactContext<C>,
next:
| ContextDependency<mixed>
| ContextDependencyWithCompare<mixed, mixed>
| null,
memoizedValue: C,
compare: (C => S) | null,
lastComparedValue: S | null,
...
lastComparedValue?: S | null,
};

export type Dependencies = {
lanes: Lanes,
firstContext: ContextDependency<mixed, mixed> | null,
firstContext:
| ContextDependency<mixed>
| ContextDependencyWithCompare<mixed, mixed>
| null,
...
};

Expand Down Expand Up @@ -388,7 +402,7 @@ export type Dispatcher = {
): [S, Dispatch<A>],
unstable_useContextWithBailout?: <T>(
context: ReactContext<T>,
compare: void | (T => mixed),
compare: (T => mixed) | null,
) => T,
useContext<T>(context: ReactContext<T>): T,
useRef<T>(initialValue: T): {current: T},
Expand Down
2 changes: 1 addition & 1 deletion packages/react/src/ReactHooks.js
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ export function useContext<T>(Context: ReactContext<T>): T {

export function unstable_useContextWithBailout<T>(
context: ReactContext<T>,
compare: void | (T => mixed),
compare: (T => mixed) | null,
): T {
if (!(enableLazyContextPropagation && enableContextProfiling)) {
throw new Error('Not implemented.');
Expand Down

0 comments on commit b3885c2

Please sign in to comment.