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.

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 = {

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'});