Throw an error when using hooks inside useMemo/useState/useReducer, or .memo's comparator (#14608)

* hooks inside useMemo/.memo - failing tests

* throw an error when using hooks inside useMemo

* throw when using hooks inside .memo's compare fn

* faster/better/stronger

* same logic for useReducer, tests for the server, etc

* Update ReactDOMServerIntegrationHooks-test.internal.js

ack lint

* nits

* whitespace

* whitespace

* stray semi

* Tweak comment

* stray unmatched fiber reset

* nit

Author: Sunil Pai

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

@@ -419,6 +419,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

@@ -234,7 +234,7 @@ export function useReducer<S, A>(
       currentHookNameInDev = 'useReducer';
     }
   }
-  currentlyRenderingComponent = resolveCurrentlyRenderingComponent();
+  let component = (currentlyRenderingComponent = resolveCurrentlyRenderingComponent());
   workInProgressHook = createWorkInProgressHook();
   if (isReRender) {
     // This is a re-render. Apply the new render phase updates to the previous
@@ -253,7 +253,10 @@ export function useReducer<S, A>(
           // priority because it will always be the same as the current
           // render's.
           const action = update.action;
+          // Temporarily clear to forbid calling Hooks.
+          currentlyRenderingComponent = null;
           newState = reducer(newState, action);
+          currentlyRenderingComponent = component;
           update = update.next;
         } while (update !== null);
 
@@ -264,6 +267,7 @@ export function useReducer<S, A>(
     }
     return [workInProgressHook.memoizedState, dispatch];
   } else {
+    currentlyRenderingComponent = null;
     if (reducer === basicStateReducer) {
       // Special case for `useState`.
       if (typeof initialState === 'function') {
@@ -272,6 +276,7 @@ export function useReducer<S, A>(
     } else if (initialAction !== undefined && initialAction !== null) {
       initialState = reducer(initialState, initialAction);
     }
+    currentlyRenderingComponent = component;
     workInProgressHook.memoizedState = initialState;
     const queue: UpdateQueue<A> = (workInProgressHook.queue = {
       last: null,
@@ -287,7 +292,7 @@ export function useReducer<S, A>(
 }
 
 function useMemo<T>(nextCreate: () => T, deps: Array<mixed> | void | null): T {
-  currentlyRenderingComponent = resolveCurrentlyRenderingComponent();
+  let component = (currentlyRenderingComponent = resolveCurrentlyRenderingComponent());
   workInProgressHook = createWorkInProgressHook();
 
   const nextDeps = deps === undefined ? null : deps;
@@ -304,7 +309,10 @@ function useMemo<T>(nextCreate: () => T, deps: Array<mixed> | void | null): T {
     }
   }
 
+  // Temporarily clear to forbid calling Hooks.
+  currentlyRenderingComponent = null;
   const nextValue = nextCreate();
+  currentlyRenderingComponent = component;
   workInProgressHook.memoizedState = [nextValue, nextDeps];
   return nextValue;
 }

packages/react-reconciler/src/ReactFiberHooks.js

@@ -420,7 +420,7 @@ export function useReducer<S, A>(
       currentHookNameInDev = 'useReducer';
     }
   }
-  currentlyRenderingFiber = resolveCurrentlyRenderingFiber();
+  let fiber = (currentlyRenderingFiber = resolveCurrentlyRenderingFiber());
   workInProgressHook = createWorkInProgressHook();
   let queue: UpdateQueue<S, A> | null = (workInProgressHook.queue: any);
   if (queue !== null) {
@@ -441,7 +441,10 @@ export function useReducer<S, A>(
             // priority because it will always be the same as the current
             // render's.
             const action = update.action;
+            // Temporarily clear to forbid calling Hooks in a reducer.
+            currentlyRenderingFiber = null;
             newState = reducer(newState, action);
+            currentlyRenderingFiber = fiber;
             update = update.next;
           } while (update !== null);
 
@@ -510,7 +513,10 @@ export function useReducer<S, A>(
             newState = ((update.eagerState: any): S);
           } else {
             const action = update.action;
+            // Temporarily clear to forbid calling Hooks in a reducer.
+            currentlyRenderingFiber = null;
             newState = reducer(newState, action);
+            currentlyRenderingFiber = fiber;
           }
         }
         prevUpdate = update;
@@ -539,7 +545,8 @@ export function useReducer<S, A>(
     const dispatch: Dispatch<A> = (queue.dispatch: any);
     return [workInProgressHook.memoizedState, dispatch];
   }
-
+  // Temporarily clear to forbid calling Hooks in a reducer.
+  currentlyRenderingFiber = null;
   // There's no existing queue, so this is the initial render.
   if (reducer === basicStateReducer) {
     // Special case for `useState`.
@@ -549,6 +556,7 @@ export function useReducer<S, A>(
   } else if (initialAction !== undefined && initialAction !== null) {
     initialState = reducer(initialState, initialAction);
   }
+  currentlyRenderingFiber = fiber;
   workInProgressHook.memoizedState = workInProgressHook.baseState = initialState;
   queue = workInProgressHook.queue = {
     last: null,
@@ -739,7 +747,7 @@ export function useMemo<T>(
   if (__DEV__) {
     currentHookNameInDev = 'useMemo';
   }
-  currentlyRenderingFiber = resolveCurrentlyRenderingFiber();
+  let fiber = (currentlyRenderingFiber = resolveCurrentlyRenderingFiber());
   workInProgressHook = createWorkInProgressHook();
 
   const nextDeps = deps === undefined ? null : deps;
@@ -755,7 +763,10 @@ export function useMemo<T>(
     }
   }
 
+  // Temporarily clear to forbid calling Hooks.
+  currentlyRenderingFiber = null;
   const nextValue = nextCreate();
+  currentlyRenderingFiber = fiber;
   workInProgressHook.memoizedState = [nextValue, nextDeps];
   return nextValue;
 }

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

@@ -517,4 +517,83 @@ 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(
+      'Hooks can only be called inside the body of a function component',
+    );
+    // the next round, it does a fresh mount, so should render
+    expect(() => root.update(<MemoApp />)).not.toThrow(
+      'Hooks can only be called inside the body of a function component',
+    );
+    // and then again, fail
+    expect(() => root.update(<MemoApp />)).toThrow(
+      'Hooks can only be called inside the body of a function component',
+    );
+  });
+
+  it('throws when calling hooks inside useMemo', () => {
+    const {useMemo, useState} = React;
+    function App() {
+      useMemo(() => {
+        useState(0);
+        return 1;
+      });
+      return null;
+    }
+
+    function Simple() {
+      const [value] = useState(123);
+      return value;
+    }
+    let root = ReactTestRenderer.create(null);
+    expect(() => root.update(<App />)).toThrow(
+      'Hooks can only be called inside the body of a function component',
+    );
+
+    // we want to assure that no hook machinery has broken
+    // so we render a fresh component with a hook just to be sure
+    root.update(<Simple />);
+    expect(root.toJSON()).toEqual('123');
+  });
+
+  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(
+      'Hooks can only be called inside the body of a function component',
+    );
+  });
+
+  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(
+      'Hooks can only be called inside the body of a function component',
+    );
+  });
 });