update lastDerivedState ref right after setDerivedState

src/__tests__/index-test.tsx

@@ -2,7 +2,7 @@
 
 import * as React from 'react';
 import * as ReactDOM from 'react-dom';
-import {Store} from 'redux';
+import {Store, createStore as createReduxStore} from 'redux';
 import {StoreContext, useMappedState} from '..';
 
 interface IAction {
@@ -214,6 +214,38 @@ describe('redux-react-hook', () => {
     expect(mapStateCalls).toBe(2);
     expect(consoleErrorSpy).not.toHaveBeenCalled();
   });
+
+  it('renders last state after synchronous dispatches', () => {
+    const store = createReduxStore(
+      (state: number = 0, action: any): number =>
+        action.type === 'test' ? action.payload : state,
+    );
+
+    const mapState = (s: number) => s;
+    const Component = () => {
+      const bar = useMappedState(mapState);
+      return <div>{bar}</div>;
+    };
+
+    render(
+      <StoreContext.Provider value={store}>
+        <Component />
+      </StoreContext.Provider>,
+    );
+
+    flushEffects();
+
+    store.dispatch({
+      type: 'test',
+      payload: 1,
+    });
+    store.dispatch({
+      type: 'test',
+      payload: 0,
+    });
+
+    expect(getText()).toBe('0');
+  });
 });
 
 // https://github.com/kentcdodds/react-testing-library/commit/11a41ce3ad9e9695f4b1662a5c67b890fc304894

src/index.ts

@@ -28,31 +28,35 @@ export function useMappedState<TState, TResult>(
   if (!store) {
     throw new Error(CONTEXT_ERROR_MESSAGE);
   }
+  const mapStateFactory = () => mapState;
   const runMapState = () => mapState(store.getState());
 
-  const [derivedState, setDerivedState] = useState(() => runMapState());
+  const [derivedState, setDerivedState] = useState(runMapState);
 
   // If the store or mapState change, rerun mapState
   const [prevStore, setPrevStore] = useState(store);
-  const [prevMapState, setPrevMapState] = useState(() => mapState);
+  const [prevMapState, setPrevMapState] = useState(mapStateFactory);
+
+  // We keep lastDerivedState in a ref and update it imperatively
+  // after calling setDerivedState so it's always up-to-date.
+  // We can't update it in useEffect because state might be updated
+  // synchronously multiple times before render occurs.
+  const lastDerivedState = useRef(derivedState);
+
+  const wrappedSetDerivedState = () => {
+    const newDerivedState = runMapState();
+    if (!shallowEqual(newDerivedState, lastDerivedState.current)) {
+      setDerivedState(newDerivedState);
+      lastDerivedState.current = newDerivedState;
+    }
+  };
+
   if (prevStore !== store || prevMapState !== mapState) {
     setPrevStore(store);
-    setPrevMapState(() => mapState);
-    setDerivedState(runMapState());
+    setPrevMapState(mapStateFactory);
+    wrappedSetDerivedState();
   }
 
-  // We use a ref to store the last result of mapState in local component
-  // state. This way we can compare with the previous version to know if
-  // the component should re-render. Otherwise, we'd have pass derivedState
-  // in the array of memoization paramaters to the second useEffect below,
-  // which would cause it to unsubscribe and resubscribe from Redux every time
-  // the state changes.
-  const lastRenderedDerivedState = useRef(derivedState);
-  // Set the last mapped state after rendering.
-  useEffect(() => {
-    lastRenderedDerivedState.current = derivedState;
-  });
-
   useEffect(
     () => {
       let didUnsubscribe = false;
@@ -66,10 +70,7 @@ export function useMappedState<TState, TResult>(
           return;
         }
 
-        const newDerivedState = runMapState();
-        if (!shallowEqual(newDerivedState, lastRenderedDerivedState.current)) {
-          setDerivedState(newDerivedState);
-        }
+        wrappedSetDerivedState();
       };
 
       // Pull data from the store after first render in case the store has