Revert "Revert "Disallow reading context during useMemo etc" (#14651)"

This reverts commit 5fce648.

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/ReactFiberDispatcher.js

@@ -7,8 +7,8 @@
  * @flow
  */
 
-import {readContext} from './ReactFiberNewContext';
 import {
+  readContext,
   useCallback,
   useContext,
   useEffect,

packages/react-reconciler/src/ReactFiberHooks.js

@@ -14,7 +14,7 @@ import type {HookEffectTag} from './ReactHookEffectTags';
 
 import {NoWork} from './ReactFiberExpirationTime';
 import {enableHooks} from 'shared/ReactFeatureFlags';
-import {readContext} from './ReactFiberNewContext';
+import {readContext as readContextWithoutCheck} from './ReactFiberNewContext';
 import {
   Update as UpdateEffect,
   Passive as PassiveEffect,
@@ -284,7 +284,7 @@ export function resetHooks(): void {
 
   // This is used to reset the state of this module when a component throws.
   // It's also called inside mountIndeterminateComponent if we determine the
-  // component is a module-style component.
+  // component is a module-style component, and also in readContext() above.
   renderExpirationTime = NoWork;
   currentlyRenderingFiber = null;
 
@@ -394,7 +394,7 @@ export function useContext<T>(
   // Ensure we're in a function component (class components support only the
   // .unstable_read() form)
   resolveCurrentlyRenderingFiber();
-  return readContext(context, observedBits);
+  return readContextWithoutCheck(context, observedBits);
 }
 
 export function useState<S>(
@@ -785,6 +785,29 @@ export function useMemo<T>(
   return nextValue;
 }
 
+export function readContext<T>(
+  context: ReactContext<T>,
+  observedBits: void | number | boolean,
+): T {
+  // Forbid reading context inside Hooks.
+  // The outer check tells us whether we're inside a Hook like useMemo().
+  // However, it would also be true if we're rendering a class.
+  if (currentlyRenderingFiber === null) {
+    // The inner check tells us we're currently in renderWithHooks() phase
+    // rather than, for example, in a class or a context consumer.
+    // Then we know it should be an error.
+    if (renderExpirationTime !== NoWork) {
+      invariant(
+        false,
+        'Context can only be read inside the body of a component. ' +
+          'If you read context inside a Hook like useMemo or useReducer, ' +
+          'move the call directly into the component body.',
+      );
+    }
+  }
+  return readContextWithoutCheck(context, observedBits);
+}
+
 function dispatchAction<S, A>(
   fiber: Fiber,
   queue: UpdateQueue<S, A>,

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

@@ -672,6 +672,88 @@ 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 inside the body of a component',
+    );
+  });
+
+  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 inside the body of a component.',
+    );
+  });
+
   it('throws when calling hooks inside useReducer', () => {
     const {useReducer, useRef} = React;
     function App() {