Disallow reading context during useMemo etc (#14653)

* Revert "Revert "Double-render function components with Hooks in DEV in StrictMode" (#14652)"
This reverts commit 3fbebb2.
* Revert "Revert "Disallow reading context during useMemo etc" (#14651)"
This reverts commit 5fce648.
* Add extra passing test for an edge case
Mentioned by @acdlite to watch out for
* More convoluted test
* Don't rely on expirationTime
Addresses @acdlite's concerns
* Edge case: forbid readContext() during eager reducer

Author: gaearon

packages/react-reconciler/src/ReactFiberClassComponent.js

@@ -20,7 +20,6 @@ import {
 importReactStrictModeWarningsfrom'./ReactStrictModeWarnings';
 import{isMounted}from'react-reconciler/reflection';
 import{getasgetInstance,setassetInstance}from'shared/ReactInstanceMap';
-importReactSharedInternalsfrom'shared/ReactSharedInternals';
 importshallowEqualfrom'shared/shallowEqual';
 importgetComponentNamefrom'shared/getComponentName';
 importinvariantfrom'shared/invariant';
@@ -48,20 +47,14 @@ import {
 hasContextChanged,
 emptyContextObject,
 }from'./ReactFiberContext';
+import{readContext}from'./ReactFiberNewContext';
 import{
 requestCurrentTime,
 computeExpirationForFiber,
 scheduleWork,
 flushPassiveEffects,
 }from'./ReactFiberScheduler';
 
-constReactCurrentDispatcher=ReactSharedInternals.ReactCurrentDispatcher;
-
-functionreadContext(contextType: any): any{
-constdispatcher=ReactCurrentDispatcher.current;
-returndispatcher.readContext(contextType);
-}
-
 constfakeInternalInstance={};
 constisArray=Array.isArray;
 

packages/react-reconciler/src/ReactFiberHooks.js

@@ -14,7 +14,11 @@ import type {HookEffectTag} from './ReactHookEffectTags';
 
 import{NoWork}from'./ReactFiberExpirationTime';
 import{enableHooks}from'shared/ReactFeatureFlags';
-import{readContext}from'./ReactFiberNewContext';
+import{
+readContext,
+stashContextDependencies,
+unstashContextDependencies,
+}from'./ReactFiberNewContext';
 import{
 UpdateasUpdateEffect,
 PassiveasPassiveEffect,
@@ -600,8 +604,10 @@ export function useReducer<S, A>(
 const action =update.action;
 // Temporarily clear to forbid calling Hooks in a reducer.
 currentlyRenderingFiber=null;
+stashContextDependencies();
 newState=reducer(newState,action);
 currentlyRenderingFiber=fiber;
+unstashContextDependencies();
 update=update.next;
 }while(update!==null);
 
@@ -672,8 +678,10 @@ export function useReducer<S, A>(
 const action =update.action;
 // Temporarily clear to forbid calling Hooks in a reducer.
 currentlyRenderingFiber=null;
+stashContextDependencies();
 newState=reducer(newState,action);
 currentlyRenderingFiber=fiber;
+unstashContextDependencies();
 }
 }
 prevUpdate=update;
@@ -704,6 +712,7 @@ export function useReducer<S, A>(
 }
 // Temporarily clear to forbid calling Hooks in a reducer.
 currentlyRenderingFiber=null;
+stashContextDependencies();
 // There's no existing queue, so this is the initial render.
 if(reducer===basicStateReducer){
 // Special case for `useState`.
@@ -714,6 +723,7 @@ export function useReducer<S, A>(
 initialState=reducer(initialState,initialAction);
 }
  currentlyRenderingFiber = fiber;
+ unstashContextDependencies();
  workInProgressHook.memoizedState = workInProgressHook.baseState = initialState;
  queue = workInProgressHook.queue = {
 last: null,
@@ -947,8 +957,10 @@ export function useMemo<T>(
 
  // Temporarily clear to forbid calling Hooks.
  currentlyRenderingFiber = null;
+ stashContextDependencies();
  const nextValue = nextCreate();
  currentlyRenderingFiber = fiber;
+ unstashContextDependencies();
  workInProgressHook.memoizedState = [nextValue, nextDeps];
  currentHookNameInDev = null;
  return nextValue;
@@ -1044,7 +1056,13 @@ function dispatchAction<S, A>(
 if(eagerReducer!==null){
 try{
 constcurrentState: S=(queue.eagerState: any);
+// Temporarily clear to forbid calling Hooks in a reducer.
+letmaybeFiber=currentlyRenderingFiber;// Note: likely null now unlike `fiber`
+currentlyRenderingFiber=null;
+stashContextDependencies();
 consteagerState=eagerReducer(currentState,action);
+currentlyRenderingFiber=maybeFiber;
+unstashContextDependencies();
 // Stash the eagerly computed state, and the reducer used to compute
 // it, on the update object. If the reducer hasn't changed by the
 // time we enter the render phase, then the eager state can be used

packages/react-reconciler/src/ReactFiberNewContext.js

@@ -52,12 +52,37 @@ let currentlyRenderingFiber: Fiber | null = null;
 let lastContextDependency: ContextDependency<mixed> | null = null;
 let lastContextWithAllBitsObserved: ReactContext<any> | null = null;
 
+// We stash the variables above before entering user code in Hooks.
+let stashedCurrentlyRenderingFiber: Fiber | null = null;
+let stashedLastContextDependency: ContextDependency<mixed> | null = null;
+let stashedLastContextWithAllBitsObserved: ReactContext<any> | null = null;
+
 export function resetContextDependences(): void {
 // This is called right before React yields execution, to ensure `readContext`
 // cannot be called outside the render phase.
 currentlyRenderingFiber=null;
 lastContextDependency=null;
 lastContextWithAllBitsObserved=null;
+
+stashedCurrentlyRenderingFiber=null;
+stashedLastContextDependency=null;
+stashedLastContextWithAllBitsObserved=null;
+}
+
+export function stashContextDependencies(): void {
+stashedCurrentlyRenderingFiber=currentlyRenderingFiber;
+stashedLastContextDependency=lastContextDependency;
+stashedLastContextWithAllBitsObserved=lastContextWithAllBitsObserved;
+
+currentlyRenderingFiber=null;
+lastContextDependency=null;
+lastContextWithAllBitsObserved=null;
+}
+
+export function unstashContextDependencies(): void {
+currentlyRenderingFiber=stashedCurrentlyRenderingFiber;
+lastContextDependency=stashedLastContextDependency;
+lastContextWithAllBitsObserved=stashedLastContextWithAllBitsObserved;
 }
 
 export function pushProvider<T>(providerFiber: Fiber, nextValue: T): void {

packages/react-reconciler/src/__tests__/ReactHooks-test.internal.js

@@ -672,6 +672,147 @@ describe('ReactHooks', () => {
 expect(root.toJSON()).toEqual('123');
 });
 
+it('throws when reading context inside useMemo',()=>{
+const{useMemo, createContext}=React;
+constReactCurrentDispatcher=
+React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+.ReactCurrentDispatcher;
+
+constThemeContext=createContext('light');
+functionApp(){
+returnuseMemo(()=>{
+returnReactCurrentDispatcher.current.readContext(ThemeContext);
+},[]);
+}
+
+expect(()=>ReactTestRenderer.create(<App/>)).toThrow(
+'Context can only be read while React is rendering',
+);
+});
+
+it('throws when reading context inside useMemo after reading outside it',()=>{
+const{useMemo, createContext}=React;
+constReactCurrentDispatcher=
+React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+.ReactCurrentDispatcher;
+
+constThemeContext=createContext('light');
+letfirstRead,secondRead;
+functionApp(){
+firstRead=ReactCurrentDispatcher.current.readContext(ThemeContext);
+useMemo(()=>{});
+secondRead=ReactCurrentDispatcher.current.readContext(ThemeContext);
+returnuseMemo(()=>{
+returnReactCurrentDispatcher.current.readContext(ThemeContext);
+},[]);
+}
+
+expect(()=>ReactTestRenderer.create(<App/>)).toThrow(
+'Context can only be read while React is rendering',
+);
+expect(firstRead).toBe('light');
+expect(secondRead).toBe('light');
+});
+
+it('throws when reading context inside useEffect',()=>{
+const{useEffect, createContext}=React;
+constReactCurrentDispatcher=
+React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+.ReactCurrentDispatcher;
+
+constThemeContext=createContext('light');
+functionApp(){
+useEffect(()=>{
+ReactCurrentDispatcher.current.readContext(ThemeContext);
+});
+returnnull;
+}
+
+constroot=ReactTestRenderer.create(<App/>);
+expect(()=>root.update(<App/>)).toThrow(
+// The exact message doesn't matter, just make sure we don't allow this
+"Cannot read property 'readContext' of null",
+);
+});
+
+it('throws when reading context inside useLayoutEffect',()=>{
+const{useLayoutEffect, createContext}=React;
+constReactCurrentDispatcher=
+React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+.ReactCurrentDispatcher;
+
+constThemeContext=createContext('light');
+functionApp(){
+useLayoutEffect(()=>{
+ReactCurrentDispatcher.current.readContext(ThemeContext);
+});
+returnnull;
+}
+
+expect(()=>ReactTestRenderer.create(<App/>)).toThrow(
+// The exact message doesn't matter, just make sure we don't allow this
+"Cannot read property 'readContext' of null",
+);
+});
+
+it('throws when reading context inside useReducer',()=>{
+const{useReducer, createContext}=React;
+constReactCurrentDispatcher=
+React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+.ReactCurrentDispatcher;
+
+constThemeContext=createContext('light');
+functionApp(){
+useReducer(
+()=>{
+ReactCurrentDispatcher.current.readContext(ThemeContext);
+},
+null,
+{},
+);
+returnnull;
+}
+
+expect(()=>ReactTestRenderer.create(<App/>)).toThrow(
+'Context can only be read while React is rendering',
+);
+});
+
+// Edge case.
+it('throws when reading context inside eager useReducer',()=>{
+const{useState, createContext}=React;
+constThemeContext=createContext('light');
+
+constReactCurrentDispatcher=
+React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+.ReactCurrentDispatcher;
+
+let_setState;
+functionFn(){
+const[,setState]=useState(0);
+_setState=setState;
+returnnull;
+}
+
+classClsextendsReact.Component{
+render(){
+_setState(()=>{
+ReactCurrentDispatcher.current.readContext(ThemeContext);
+});
+returnnull;
+}
+}
+
+expect(()=>
+ReactTestRenderer.create(
+<React.Fragment>
+<Fn/>
+<Cls/>
+</React.Fragment>,
+),
+).toThrow('Context can only be read while React is rendering');
+});
+
 it('throws when calling hooks inside useReducer',()=>{
 const{useReducer, useRef}=React;
 functionApp(){