Throw an error when using hooks inside useMemo, or .memo's comparison function

packages/react-dom/src/__tests__/ReactDOMServerIntegrationHooks-test.internal.js

@@ -417,6 +417,51 @@ describe('ReactDOMServerHooks', () => {
         expect(domNode.textContent).toEqual('HELLO, WORLD.');
       },
     );
+
+    itThrowsWhenRendering(
+      'a hook inside useMemo',
+      async render => {
+        function App() {
+          useMemo(() => {
+            useState();
+            return 0;
+          });
+          return null;
+        }
+        return render(<App />);
+      },
+      'Hooks can only be called inside the body of a function component.',
+    );
+
+    itThrowsWhenRendering(
+      'a hook inside useReducer',
+      async render => {
+        function App() {
+          const [value, dispatch] = useReducer((state, action) => {
+            useRef(0);
+            return state;
+          }, 0);
+          dispatch('foo');
+          return value;
+        }
+        return render(<App />);
+      },
+      'Hooks can only be called inside the body of a function component.',
+    );
+
+    itThrowsWhenRendering(
+      'a hook inside useState',
+      async render => {
+        function App() {
+          useState(() => {
+            useRef(0);
+            return 0;
+          });
+        }
+        return render(<App />);
+      },
+      'Hooks can only be called inside the body of a function component.',
+    );
   });
 
   describe('useRef', () => {

packages/react-dom/src/server/ReactPartialRendererHooks.js

@@ -177,7 +177,7 @@ export function useReducer<S, A>(
   initialState: S,
   initialAction: A | void | null,
 ): [S, Dispatch<A>] {
-  currentlyRenderingComponent = resolveCurrentlyRenderingComponent();
+  let cachedCurrentlyRenderingComponent = (currentlyRenderingComponent = resolveCurrentlyRenderingComponent());
   workInProgressHook = createWorkInProgressHook();
   if (isReRender) {
     // This is a re-render. Apply the new render phase updates to the previous
@@ -196,7 +196,12 @@ export function useReducer<S, A>(
           // priority because it will always be the same as the current
           // render's.
           const action = update.action;
+
+          currentlyRenderingComponent = null;
+          // now, it'll throw if you try to call a hook inside the reducer
           newState = reducer(newState, action);
+          // restore the current component
+          currentlyRenderingComponent = cachedCurrentlyRenderingComponent;
           update = update.next;
         } while (update !== null);
 
@@ -207,6 +212,7 @@ export function useReducer<S, A>(
     }
     return [workInProgressHook.memoizedState, dispatch];
   } else {
+    currentlyRenderingComponent = null;
     if (reducer === basicStateReducer) {
       // Special case for `useState`.
       if (typeof initialState === 'function') {
@@ -215,6 +221,7 @@ export function useReducer<S, A>(
     } else if (initialAction !== undefined && initialAction !== null) {
       initialState = reducer(initialState, initialAction);
     }
+    currentlyRenderingComponent = cachedCurrentlyRenderingComponent;
     workInProgressHook.memoizedState = initialState;
     const queue: UpdateQueue<A> = (workInProgressHook.queue = {
       last: null,
@@ -233,7 +240,7 @@ function useMemo<T>(
   nextCreate: () => T,
   inputs: Array<mixed> | void | null,
 ): T {
-  currentlyRenderingComponent = resolveCurrentlyRenderingComponent();
+  let cachedCurrentlyRenderingComponent = (currentlyRenderingComponent = resolveCurrentlyRenderingComponent());
   workInProgressHook = createWorkInProgressHook();
 
   const nextInputs =
@@ -250,7 +257,12 @@ function useMemo<T>(
     }
   }
 
+  currentlyRenderingComponent = null;
+  // now, it'll throw if you try to call a hook inside nextCreate()
   const nextValue = nextCreate();
+  // restore the current component
+  currentlyRenderingComponent = cachedCurrentlyRenderingComponent;
+
   workInProgressHook.memoizedState = [nextValue, nextInputs];
   return nextValue;
 }

packages/react-reconciler/src/ReactFiberHooks.js

@@ -108,7 +108,6 @@ let renderPhaseUpdates: Map<UpdateQueue<any>, Update<any>> | null = null;
 // Counter to prevent infinite loops.
 let numberOfReRenders: number = 0;
 const RE_RENDER_LIMIT = 25;
-
 function resolveCurrentlyRenderingFiber(): Fiber {
   invariant(
     currentlyRenderingFiber !== null,
@@ -345,7 +344,7 @@ export function useReducer<S, A>(
   initialState: S,
   initialAction: A | void | null,
 ): [S, Dispatch<A>] {
-  currentlyRenderingFiber = resolveCurrentlyRenderingFiber();
+  let cachedCurrentlyRenderingFiber = (currentlyRenderingFiber = resolveCurrentlyRenderingFiber());
   workInProgressHook = createWorkInProgressHook();
   let queue: UpdateQueue<A> | null = (workInProgressHook.queue: any);
   if (queue !== null) {
@@ -366,7 +365,12 @@ export function useReducer<S, A>(
             // priority because it will always be the same as the current
             // render's.
             const action = update.action;
+            currentlyRenderingFiber = null;
+
+            // now, it'll throw if you try to call a hook inside the reducer
             newState = reducer(newState, action);
+            // restore the current component
+            currentlyRenderingFiber = cachedCurrentlyRenderingFiber;
             update = update.next;
           } while (update !== null);
 
@@ -429,7 +433,9 @@ export function useReducer<S, A>(
         } else {
           // Process this update.
           const action = update.action;
+          currentlyRenderingFiber = null;
           newState = reducer(newState, action);
+          currentlyRenderingFiber = cachedCurrentlyRenderingFiber;
         }
         prevUpdate = update;
         update = update.next;
@@ -448,7 +454,7 @@ export function useReducer<S, A>(
     const dispatch: Dispatch<A> = (queue.dispatch: any);
     return [workInProgressHook.memoizedState, dispatch];
   }
-
+  currentlyRenderingFiber = null;
   // There's no existing queue, so this is the initial render.
   if (reducer === basicStateReducer) {
     // Special case for `useState`.
@@ -458,6 +464,7 @@ export function useReducer<S, A>(
   } else if (initialAction !== undefined && initialAction !== null) {
     initialState = reducer(initialState, initialAction);
   }
+  currentlyRenderingFiber = cachedCurrentlyRenderingFiber;
   workInProgressHook.memoizedState = workInProgressHook.baseState = initialState;
   queue = workInProgressHook.queue = {
     last: null,
@@ -613,7 +620,7 @@ export function useMemo<T>(
   nextCreate: () => T,
   inputs: Array<mixed> | void | null,
 ): T {
-  currentlyRenderingFiber = resolveCurrentlyRenderingFiber();
+  let cachedCurrentlyRenderingFiber = (currentlyRenderingFiber = resolveCurrentlyRenderingFiber());
   workInProgressHook = createWorkInProgressHook();
 
   const nextInputs =
@@ -627,7 +634,12 @@ export function useMemo<T>(
     }
   }
 
+  currentlyRenderingFiber = null;
+  // now, it'll throw if you try to call a hook inside nextCreate
   const nextValue = nextCreate();
+  // restore the current fiber
+  currentlyRenderingFiber = cachedCurrentlyRenderingFiber;
+
   workInProgressHook.memoizedState = [nextValue, nextInputs];
   return nextValue;
 }

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

@@ -17,6 +17,9 @@ let ReactFeatureFlags;
 let ReactTestRenderer;
 let ReactDOMServer;
 
+const OutsideHookScopeError =
+  'Hooks can only be called inside the body of a function component';
+
 // Additional tests can be found in ReactHooksWithNoopRenderer. Plan is to
 // gradually migrate those to this file.
 describe('ReactHooks', () => {
@@ -92,4 +95,73 @@ describe('ReactHooks', () => {
     const root = ReactTestRenderer.create(<App />);
     expect(root.toJSON()).toMatchSnapshot();
   });
+
+  it("throws when calling hooks inside .memo's compare function", () => {
+    const {useState} = React;
+    function App() {
+      useState(0);
+      return null;
+    }
+    const MemoApp = React.memo(App, () => {
+      useState(0);
+      return false;
+    });
+
+    const root = ReactTestRenderer.create(<MemoApp />);
+    // trying to render again should trigger comparison and throw
+    expect(() => root.update(<MemoApp />)).toThrow(OutsideHookScopeError);
+    // the next round, it does a fresh mount, so should render
+    expect(() => root.update(<MemoApp />)).not.toThrow(OutsideHookScopeError);
+    // and then again, fail
+    expect(() => root.update(<MemoApp />)).toThrow(OutsideHookScopeError);
+  });
+
+  it('throws when calling hooks inside useMemo', () => {
+    const {useMemo, useState} = React;
+    function App() {
+      useMemo(() => {
+        useState(0);
+        return 123;
+      });
+      return null;
+    }
+
+    function Simple() {
+      useState(0);
+      return 123;
+    }
+    let root = ReactTestRenderer.create(null);
+    expect(() => root.update(<App />)).toThrow(OutsideHookScopeError);
+    // we want to assure that no hook machinery has broken,
+    // so we listen for all errors on the next line
+    expect(() => root.update(<Simple />)).not.toThrow();
+  });
+
+  it('throws when calling hooks inside useReducer', () => {
+    const {useReducer, useRef} = React;
+    function App() {
+      const [value, dispatch] = useReducer((state, action) => {
+        useRef(0);
+        return state;
+      }, 0);
+      dispatch('foo');
+      return value;
+    }
+    expect(() => ReactTestRenderer.create(<App />)).toThrow(
+      OutsideHookScopeError,
+    );
+  });
+
+  it("throws when calling hooks inside useState's initialize function", () => {
+    const {useState, useRef} = React;
+    function App() {
+      useState(() => {
+        useRef(0);
+        return 0;
+      });
+    }
+    expect(() => ReactTestRenderer.create(<App />)).toThrow(
+      OutsideHookScopeError,
+    );
+  });
 });