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

packages/react-reconciler/src/ReactFiberClassComponent.js

@@ -20,7 +20,6 @@ import {
 import ReactStrictModeWarnings from './ReactStrictModeWarnings';
 import {isMounted} from 'react-reconciler/reflection';
 import {get as getInstance, set as setInstance} from 'shared/ReactInstanceMap';
-import ReactSharedInternals from 'shared/ReactSharedInternals';
 import shallowEqual from 'shared/shallowEqual';
 import getComponentName from 'shared/getComponentName';
 import invariant from 'shared/invariant';
@@ -48,20 +47,14 @@ import {
   hasContextChanged,
   emptyContextObject,
 } from './ReactFiberContext';
+import {readContext} from './ReactFiberNewContext';
 import {
   requestCurrentTime,
   computeExpirationForFiber,
   scheduleWork,
   flushPassiveEffects,
 } from './ReactFiberScheduler';
 
-const ReactCurrentDispatcher = ReactSharedInternals.ReactCurrentDispatcher;
-
-function readContext(contextType: any): any {
-  const dispatcher = ReactCurrentDispatcher.current;
-  return dispatcher.readContext(contextType);
-}
-
 const fakeInternalInstance = {};
 const isArray = 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 {
   Update as UpdateEffect,
   Passive as PassiveEffect,
@@ -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 {
           const currentState: S = (queue.eagerState: any);
+          // Temporarily clear to forbid calling Hooks in a reducer.
+          let maybeFiber = currentlyRenderingFiber; // Note: likely null now unlike `fiber`
+          currentlyRenderingFiber = null;
+          stashContextDependencies();
           const eagerState = 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;
+    const ReactCurrentDispatcher =
+      React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .ReactCurrentDispatcher;
+
+    const ThemeContext = createContext('light');
+    function App() {
+      return useMemo(() => {
+        return ReactCurrentDispatcher.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;
+    const ReactCurrentDispatcher =
+      React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .ReactCurrentDispatcher;
+
+    const ThemeContext = createContext('light');
+    let firstRead, secondRead;
+    function App() {
+      firstRead = ReactCurrentDispatcher.current.readContext(ThemeContext);
+      useMemo(() => {});
+      secondRead = ReactCurrentDispatcher.current.readContext(ThemeContext);
+      return useMemo(() => {
+        return ReactCurrentDispatcher.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;
+    const ReactCurrentDispatcher =
+      React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .ReactCurrentDispatcher;
+
+    const ThemeContext = createContext('light');
+    function App() {
+      useEffect(() => {
+        ReactCurrentDispatcher.current.readContext(ThemeContext);
+      });
+      return null;
+    }
+
+    const root = 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;
+    const ReactCurrentDispatcher =
+      React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .ReactCurrentDispatcher;
+
+    const ThemeContext = createContext('light');
+    function App() {
+      useLayoutEffect(() => {
+        ReactCurrentDispatcher.current.readContext(ThemeContext);
+      });
+      return null;
+    }
+
+    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;
+    const ReactCurrentDispatcher =
+      React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .ReactCurrentDispatcher;
+
+    const ThemeContext = createContext('light');
+    function App() {
+      useReducer(
+        () => {
+          ReactCurrentDispatcher.current.readContext(ThemeContext);
+        },
+        null,
+        {},
+      );
+      return null;
+    }
+
+    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;
+    const ThemeContext = createContext('light');
+
+    const ReactCurrentDispatcher =
+      React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .ReactCurrentDispatcher;
+
+    let _setState;
+    function Fn() {
+      const [, setState] = useState(0);
+      _setState = setState;
+      return null;
+    }
+
+    class Cls extends React.Component {
+      render() {
+        _setState(() => {
+          ReactCurrentDispatcher.current.readContext(ThemeContext);
+        });
+        return null;
+      }
+    }
+
+    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;
     function App() {