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',
+asyncrender=>{
+functionApp(){
+useMemo(()=>{
+useState();
+return0;
+});
+returnnull;
+}
+returnrender(<App/>);
+},
+'Hooks can only be called inside the body of a function component.',
+);
+
+itThrowsWhenRendering(
+'a hook inside useReducer',
+asyncrender=>{
+functionApp(){
+const[value,dispatch]=useReducer((state,action)=>{
+useRef(0);
+returnstate;
+},0);
+dispatch('foo');
+returnvalue;
+}
+returnrender(<App/>);
+},
+'Hooks can only be called inside the body of a function component.',
+);
+
+itThrowsWhenRendering(
+'a hook inside useState',
+asyncrender=>{
+functionApp(){
+useState(()=>{
+useRef(0);
+return0;
+});
+}
+returnrender(<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();
+letcomponent=(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(typeofinitialState==='function'){
@@ -272,6 +276,7 @@ export function useReducer<S, A>(
 }elseif(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>(
 }
 
 functionuseMemo<T>(nextCreate: () =>T,deps: Array<mixed> | void | null): T {
-currentlyRenderingComponent=resolveCurrentlyRenderingComponent();
+letcomponent=(currentlyRenderingComponent=resolveCurrentlyRenderingComponent());
 workInProgressHook=createWorkInProgressHook();
 
 constnextDeps=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;
 constnextValue=nextCreate();
+currentlyRenderingComponent=component;
 workInProgressHook.memoizedState=[nextValue,nextDeps];
 returnnextValue;
 }

packages/react-reconciler/src/ReactFiberHooks.js

@@ -420,7 +420,7 @@ export function useReducer<S, A>(
 currentHookNameInDev='useReducer';
 }
 }
-currentlyRenderingFiber=resolveCurrentlyRenderingFiber();
+letfiber=(currentlyRenderingFiber=resolveCurrentlyRenderingFiber());
 workInProgressHook=createWorkInProgressHook();
 letqueue: 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>(
 constdispatch: 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>(
 }elseif(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;
 constnextValue=nextCreate();
+currentlyRenderingFiber=fiber;
 workInProgressHook.memoizedState=[nextValue,nextDeps];
 returnnextValue;
 }

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

@@ -517,4 +517,83 @@ describe('ReactHooks', () => {
 constroot=ReactTestRenderer.create(<App/>);
 expect(root.toJSON()).toMatchSnapshot();
 });
+
+it("throws when calling hooks inside .memo's compare function",()=>{
+const{useState}=React;
+functionApp(){
+useState(0);
+returnnull;
+}
+constMemoApp=React.memo(App,()=>{
+useState(0);
+returnfalse;
+});
+
+constroot=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;
+functionApp(){
+useMemo(()=>{
+useState(0);
+return1;
+});
+returnnull;
+}
+
+functionSimple(){
+const[value]=useState(123);
+returnvalue;
+}
+letroot=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;
+functionApp(){
+const[value,dispatch]=useReducer((state,action)=>{
+useRef(0);
+returnstate;
+},0);
+dispatch('foo');
+returnvalue;
+}
+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;
+functionApp(){
+useState(()=>{
+useRef(0);
+return0;
+});
+}
+expect(()=>ReactTestRenderer.create(<App/>)).toThrow(
+'Hooks can only be called inside the body of a function component',
+);
+});
 });