Fix suspense in useState and useReducer initializers (#41)

* Add failing "suspense in state initialiser" test

* Fix operation in useReducer hook

- Apply initial state if hook is uninitialised
- Ensure that queue/dispatch are initialised before applying updates

Author: kitten

src/__tests__/suspense.test.js

@@ -224,6 +224,33 @@ describe('renderPrepass', () => {
       })
     })
 
+    it("handles suspenses thrown in useState's initialState initialiser", () => {
+      const getValue = jest
+        .fn()
+        .mockImplementationOnce(() => {
+          throw Promise.resolve()
+        })
+        .mockImplementation(() => 'test')
+
+      const Inner = jest.fn(props => {
+        expect(props.state).toBe('test')
+      })
+
+      const Outer = jest.fn(() => {
+        const [state] = useState(() => {
+          // This will trigger suspense:
+          return getValue()
+        })
+
+        return <Inner state={state} />
+      })
+
+      return renderPrepass(<Outer />).then(() => {
+        expect(Outer).toHaveBeenCalled()
+        expect(Inner).toHaveBeenCalled()
+      })
+    })
+
     it('ignores thrown non-promises', () => {
       const Outer = () => {
         throw new Error('test')
@@ -435,7 +462,7 @@ describe('renderPrepass', () => {
       })
     })
 
-    it('supports rendering previouslt resolved lazy components', () => {
+    it('supports rendering previously resolved lazy components', () => {
       const Inner = jest.fn(() => null)
       const loadInner = jest.fn().mockResolvedValueOnce(Inner)
       const Outer = React.lazy(loadInner)

src/internals/dispatcher.js

@@ -175,34 +175,10 @@ function useReducer<S, I, A>(
   const id = getCurrentIdentity()
   workInProgressHook = createWorkInProgressHook()
 
-  if (isReRender) {
-    // This is a re-render. Apply the new render phase updates to the previous
-    // current hook.
-    const queue: UpdateQueue<A> = (workInProgressHook.queue: any)
-    const dispatch: Dispatch<A> = (queue.dispatch: any)
-    if (renderPhaseUpdates !== null) {
-      // Render phase updates are stored in a map of queue -> linked list
-      const firstRenderPhaseUpdate = renderPhaseUpdates.get(queue)
-      if (firstRenderPhaseUpdate !== undefined) {
-        renderPhaseUpdates.delete(queue)
-        let newState = workInProgressHook.memoizedState
-        let update = firstRenderPhaseUpdate
-        do {
-          // Process this render phase update. We don't have to check the
-          // priority because it will always be the same as the current
-          // render's.
-          const action = update.action
-          newState = reducer(newState, action)
-          update = update.next
-        } while (update !== null)
-
-        workInProgressHook.memoizedState = newState
-
-        return [newState, dispatch]
-      }
-    }
-    return [workInProgressHook.memoizedState, dispatch]
-  } else {
+  // In the case of a re-render after a suspense, the initial state
+  // may not be set, so instead of initialising if `!isRerender`, we
+  // check whether `queue` is set
+  if (workInProgressHook.queue === null) {
     let initialState
     if (reducer === basicStateReducer) {
       // Special case for `useState`.
@@ -214,18 +190,39 @@ function useReducer<S, I, A>(
       initialState =
         init !== undefined ? init(initialArg) : ((initialArg: any): S)
     }
+
     workInProgressHook.memoizedState = initialState
-    const queue: UpdateQueue<A> = (workInProgressHook.queue = {
-      last: null,
-      dispatch: null
-    })
-    const dispatch: Dispatch<A> = (queue.dispatch = (dispatchAction.bind(
-      null,
-      id,
-      queue
-    ): any))
-    return [workInProgressHook.memoizedState, dispatch]
   }
+
+  const queue: UpdateQueue<A> =
+    workInProgressHook.queue ||
+    (workInProgressHook.queue = { last: null, dispatch: null })
+  const dispatch: Dispatch<A> =
+    queue.dispatch || (queue.dispatch = dispatchAction.bind(null, id, queue))
+
+  if (isReRender && renderPhaseUpdates !== null) {
+    // This is a re-render. Apply the new render phase updates to the previous
+    // current hook.
+    // Render phase updates are stored in a map of queue -> linked list
+    const firstRenderPhaseUpdate = renderPhaseUpdates.get(queue)
+    if (firstRenderPhaseUpdate !== undefined) {
+      renderPhaseUpdates.delete(queue)
+      let newState = workInProgressHook.memoizedState
+      let update = firstRenderPhaseUpdate
+      do {
+        // Process this render phase update. We don't have to check the
+        // priority because it will always be the same as the current
+        // render's.
+        const action = update.action
+        newState = reducer(newState, action)
+        update = update.next
+      } while (update !== null)
+
+      workInProgressHook.memoizedState = newState
+    }
+  }
+
+  return [workInProgressHook.memoizedState, dispatch]
 }
 
 function useMemo<T>(nextCreate: () => T, deps: Array<mixed> | void | null): T {