Improve Reducer Hook's lazy init API (#14723)

* Improve Reducer Hook's lazy init API

* Use generic type for initilizer input

Still requires an `any` cast in the case where `init` function is
not provided.

Author: acdlite

packages/react-debug-tools/src/ReactDebugHooks.js

@@ -115,13 +115,18 @@ function useState<S>(
   return [state, (action: BasicStateAction<S>) => {}];
 }
 
-function useReducer<S, A>(
+function useReducer<S, I, A>(
   reducer: (S, A) => S,
-  initialState: S,
-  initialAction: A | void | null,
+  initialArg: I,
+  init?: I => S,
 ): [S, Dispatch<A>] {
   let hook = nextHook();
-  let state = hook !== null ? hook.memoizedState : initialState;
+  let state;
+  if (hook !== null) {
+    state = hook.memoizedState;
+  } else {
+    state = init !== undefined ? init(initialArg) : ((initialArg: any): S);
+  }
   hookLog.push({
     primitive: 'Reducer',
     stackError: new Error(),

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

@@ -208,12 +208,12 @@ describe('ReactDOMServerHooks', () => {
       expect(domNode.textContent).toEqual('0');
     });
 
-    itRenders('lazy initialization with initialAction', async render => {
+    itRenders('lazy initialization', async render => {
       function reducer(state, action) {
         return action === 'increment' ? state + 1 : state;
       }
       function Counter() {
-        let [count] = useReducer(reducer, 0, 'increment');
+        let [count] = useReducer(reducer, 0, c => c + 1);
         yieldValue('Render: ' + count);
         return <Text text={count} />;
       }

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

@@ -252,10 +252,10 @@ export function useState<S>(
   );
 }
 
-export function useReducer<S, A>(
+export function useReducer<S, I, A>(
   reducer: (S, A) => S,
-  initialState: S,
-  initialAction: A | void | null,
+  initialArg: I,
+  init?: I => S,
 ): [S, Dispatch<A>] {
   if (__DEV__) {
     if (reducer !== basicStateReducer) {
@@ -301,13 +301,16 @@ export function useReducer<S, A>(
     if (__DEV__) {
       isInHookUserCodeInDev = true;
     }
+    let initialState;
     if (reducer === basicStateReducer) {
       // Special case for `useState`.
-      if (typeof initialState === 'function') {
-        initialState = initialState();
-      }
-    } else if (initialAction !== undefined && initialAction !== null) {
-      initialState = reducer(initialState, initialAction);
+      initialState =
+        typeof initialArg === 'function'
+          ? ((initialArg: any): () => S)()
+          : ((initialArg: any): S);
+    } else {
+      initialState =
+        init !== undefined ? init(initialArg) : ((initialArg: any): S);
     }
     if (__DEV__) {
       isInHookUserCodeInDev = false;

packages/react-reconciler/src/ReactFiberHooks.js

@@ -49,10 +49,10 @@ export type Dispatcher = {
     observedBits: void | number | boolean,
   ): T,
   useState<S>(initialState: (() => S) | S): [S, Dispatch<BasicStateAction<S>>],
-  useReducer<S, A>(
+  useReducer<S, I, A>(
     reducer: (S, A) => S,
-    initialState: S,
-    initialAction: A | void | null,
+    initialArg: I,
+    init?: (I) => S,
   ): [S, Dispatch<A>],
   useContext<T>(
     context: ReactContext<T>,
@@ -591,16 +591,17 @@ function updateContext<T>(
   return readContext(context, observedBits);
 }
 
-function mountReducer<S, A>(
+function mountReducer<S, I, A>(
   reducer: (S, A) => S,
-  initialState: void | S,
-  initialAction: void | null | A,
+  initialArg: I,
+  init?: I => S,
 ): [S, Dispatch<A>] {
   const hook = mountWorkInProgressHook();
-  // TODO: Lazy init API will change before release.
-  if (initialAction !== undefined && initialAction !== null) {
-    // $FlowFixMe - Must express with overloading.
-    initialState = reducer(initialState, initialAction);
+  let initialState;
+  if (init !== undefined) {
+    initialState = init(initialArg);
+  } else {
+    initialState = ((initialArg: any): S);
   }
   hook.memoizedState = hook.baseState = initialState;
   const queue = (hook.queue = {
@@ -618,10 +619,10 @@ function mountReducer<S, A>(
   return [hook.memoizedState, dispatch];
 }
 
-function updateReducer<S, A>(
+function updateReducer<S, I, A>(
   reducer: (S, A) => S,
-  initialState: void | S,
-  initialAction: void | null | A,
+  initialArg: I,
+  init?: I => S,
 ): [S, Dispatch<A>] {
   const hook = updateWorkInProgressHook();
   const queue = hook.queue;
@@ -755,7 +756,6 @@ function mountState<S>(
   initialState: (() => S) | S,
 ): [S, Dispatch<BasicStateAction<S>>] {
   const hook = mountWorkInProgressHook();
-  // TODO: Lazy init API will change before release.
   if (typeof initialState === 'function') {
     initialState = initialState();
   }
@@ -1282,16 +1282,16 @@ if (__DEV__) {
         ReactCurrentDispatcher.current = prevDispatcher;
       }
     },
-    useReducer<S, A>(
+    useReducer<S, I, A>(
       reducer: (S, A) => S,
-      initialState: S,
-      initialAction: A | void | null,
+      initialArg: I,
+      init?: I => S,
     ): [S, Dispatch<A>] {
       currentHookNameInDev = 'useReducer';
       const prevDispatcher = ReactCurrentDispatcher.current;
       ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnMountInDEV;
       try {
-        return mountReducer(reducer, initialState, initialAction);
+        return mountReducer(reducer, initialArg, init);
       } finally {
         ReactCurrentDispatcher.current = prevDispatcher;
       }
@@ -1366,16 +1366,16 @@ if (__DEV__) {
         ReactCurrentDispatcher.current = prevDispatcher;
       }
     },
-    useReducer<S, A>(
+    useReducer<S, I, A>(
       reducer: (S, A) => S,
-      initialState: S,
-      initialAction: A | void | null,
+      initialArg: I,
+      init?: I => S,
     ): [S, Dispatch<A>] {
       currentHookNameInDev = 'useReducer';
       const prevDispatcher = ReactCurrentDispatcher.current;
       ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnUpdateInDEV;
       try {
-        return updateReducer(reducer, initialState, initialAction);
+        return updateReducer(reducer, initialArg, init);
       } finally {
         ReactCurrentDispatcher.current = prevDispatcher;
       }
@@ -1457,17 +1457,17 @@ if (__DEV__) {
         ReactCurrentDispatcher.current = prevDispatcher;
       }
     },
-    useReducer<S, A>(
+    useReducer<S, I, A>(
       reducer: (S, A) => S,
-      initialState: S,
-      initialAction: A | void | null,
+      initialArg: I,
+      init?: I => S,
     ): [S, Dispatch<A>] {
       currentHookNameInDev = 'useReducer';
       warnInvalidHookAccess();
       const prevDispatcher = ReactCurrentDispatcher.current;
       ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnMountInDEV;
       try {
-        return mountReducer(reducer, initialState, initialAction);
+        return mountReducer(reducer, initialArg, init);
       } finally {
         ReactCurrentDispatcher.current = prevDispatcher;
       }
@@ -1552,17 +1552,17 @@ if (__DEV__) {
         ReactCurrentDispatcher.current = prevDispatcher;
       }
     },
-    useReducer<S, A>(
+    useReducer<S, I, A>(
       reducer: (S, A) => S,
-      initialState: S,
-      initialAction: A | void | null,
+      initialArg: I,
+      init?: I => S,
     ): [S, Dispatch<A>] {
       currentHookNameInDev = 'useReducer';
       warnInvalidHookAccess();
       const prevDispatcher = ReactCurrentDispatcher.current;
       ReactCurrentDispatcher.current = InvalidNestedHooksDispatcherOnUpdateInDEV;
       try {
-        return updateReducer(reducer, initialState, initialAction);
+        return updateReducer(reducer, initialArg, init);
       } finally {
         ReactCurrentDispatcher.current = prevDispatcher;
       }

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

@@ -752,19 +752,19 @@ describe('ReactHooks', () => {
 
     const ThemeContext = createContext('light');
     function App() {
-      useReducer(
-        () => {
-          ReactCurrentDispatcher.current.readContext(ThemeContext);
-        },
-        null,
-        {},
-      );
+      const [state, dispatch] = useReducer((s, action) => {
+        ReactCurrentDispatcher.current.readContext(ThemeContext);
+        return action;
+      }, 0);
+      if (state === 0) {
+        dispatch(1);
+      }
       return null;
     }
 
-    expect(() => ReactTestRenderer.create(<App />)).toWarnDev(
+    expect(() => ReactTestRenderer.create(<App />)).toWarnDev([
       'Context can only be read while React is rendering',
-    );
+    ]);
   });
 
   // Edge case.

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

@@ -524,14 +524,12 @@ describe('ReactHooksWithNoopRenderer', () => {
       expect(ReactNoop.getChildren()).toEqual([span('Count: -2')]);
     });
 
-    it('accepts an initial action', () => {
+    it('lazy init', () => {
       const INCREMENT = 'INCREMENT';
       const DECREMENT = 'DECREMENT';
 
       function reducer(state, action) {
         switch (action) {
-          case 'INITIALIZE':
-            return 10;
           case 'INCREMENT':
             return state + 1;
           case 'DECREMENT':
@@ -541,27 +539,28 @@ describe('ReactHooksWithNoopRenderer', () => {
         }
       }
 
-      const initialAction = 'INITIALIZE';
-
       function Counter(props, ref) {
-        const [count, dispatch] = useReducer(reducer, 0, initialAction);
+        const [count, dispatch] = useReducer(reducer, props, p => {
+          ReactNoop.yield('Init');
+          return p.initialCount;
+        });
         useImperativeHandle(ref, () => ({dispatch}));
         return <Text text={'Count: ' + count} />;
       }
       Counter = forwardRef(Counter);
       const counter = React.createRef(null);
-      ReactNoop.render(<Counter ref={counter} />);
-      ReactNoop.flush();
+      ReactNoop.render(<Counter initialCount={10} ref={counter} />);
+      expect(ReactNoop.flush()).toEqual(['Init', 'Count: 10']);
       expect(ReactNoop.getChildren()).toEqual([span('Count: 10')]);
 
       counter.current.dispatch(INCREMENT);
-      ReactNoop.flush();
+      expect(ReactNoop.flush()).toEqual(['Count: 11']);
       expect(ReactNoop.getChildren()).toEqual([span('Count: 11')]);
 
       counter.current.dispatch(DECREMENT);
       counter.current.dispatch(DECREMENT);
       counter.current.dispatch(DECREMENT);
-      ReactNoop.flush();
+      expect(ReactNoop.flush()).toEqual(['Count: 8']);
       expect(ReactNoop.getChildren()).toEqual([span('Count: 8')]);
     });
 

packages/react-test-renderer/src/ReactShallowRenderer.js

@@ -223,10 +223,10 @@ class ReactShallowRenderer {
   }
 
   _createDispatcher(): DispatcherType {
-    const useReducer = <S, A>(
+    const useReducer = <S, I, A>(
       reducer: (S, A) => S,
-      initialState: S,
-      initialAction: A | void | null,
+      initialArg: I,
+      init?: I => S,
     ): [S, Dispatch<A>] => {
       this._validateCurrentlyRenderingComponent();
       this._createWorkInProgressHook();
@@ -259,13 +259,16 @@ class ReactShallowRenderer {
         }
         return [workInProgressHook.memoizedState, dispatch];
       } else {
+        let initialState;
         if (reducer === basicStateReducer) {
           // Special case for `useState`.
-          if (typeof initialState === 'function') {
-            initialState = initialState();
-          }
-        } else if (initialAction !== undefined && initialAction !== null) {
-          initialState = reducer(initialState, initialAction);
+          initialState =
+            typeof initialArg === 'function'
+              ? ((initialArg: any): () => S)()
+              : ((initialArg: any): S);
+        } else {
+          initialState =
+            init !== undefined ? init(initialArg) : ((initialArg: any): S);
         }
         workInProgressHook.memoizedState = initialState;
         const queue: UpdateQueue<A> = (workInProgressHook.queue = {

packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.js

@@ -91,23 +91,19 @@ describe('ReactShallowRenderer with hooks', () => {
   });
 
   it('should work with useReducer', () => {
-    const initialState = {count: 0};
-
     function reducer(state, action) {
       switch (action.type) {
-        case 'reset':
-          return initialState;
         case 'increment':
           return {count: state.count + 1};
         case 'decrement':
           return {count: state.count - 1};
-        default:
-          return state;
       }
     }
 
-    function SomeComponent({initialCount}) {
-      const [state] = React.useReducer(reducer, {count: initialCount});
+    function SomeComponent(props) {
+      const [state] = React.useReducer(reducer, props, p => ({
+        count: p.initialCount,
+      }));
 
       return (
         <div>
@@ -141,25 +137,19 @@ describe('ReactShallowRenderer with hooks', () => {
   });
 
   it('should work with a dispatched state change for a useReducer', () => {
-    const initialState = {count: 0};
-
     function reducer(state, action) {
       switch (action.type) {
-        case 'reset':
-          return initialState;
         case 'increment':
           return {count: state.count + 1};
         case 'decrement':
           return {count: state.count - 1};
-        default:
-          return state;
       }
     }
 
-    function SomeComponent({initialCount}) {
-      const [state, dispatch] = React.useReducer(reducer, {
-        count: initialCount,
-      });
+    function SomeComponent(props) {
+      const [state, dispatch] = React.useReducer(reducer, props, p => ({
+        count: p.initialCount,
+      }));
 
       if (state.count === 0) {
         dispatch({type: 'increment'});