Simple usemappedstate

src/__tests__/index-test.tsx

@@ -214,7 +214,7 @@ describe('redux-react-hook', () => {
       updateStore({...state, foo: 'foo'});
 
       // mapState is called during and after the first render
-      expect(mapStateCalls).toBe(2);
+      expect(mapStateCalls).toBe(1);
       expect(consoleErrorSpy).not.toHaveBeenCalled();
     });
 
@@ -255,9 +255,7 @@ describe('redux-react-hook', () => {
       const Component = ({prop}: {prop: any}) => {
         const mapState = React.useCallback((s: IState) => s, [prop]);
         useMappedState(mapState);
-        React.useEffect(() => {
-          renderCount++;
-        });
+        renderCount++;
         return null;
       };
 
@@ -279,6 +277,47 @@ describe('redux-react-hook', () => {
         ReactDOM.render(<Component />, reactRoot);
       });
     });
+
+    it('does not provide stale mapped state', () => {
+      let flag = false;
+
+      const Component = ({prop}: {prop: any}) => {
+        const mapState = React.useCallback((s: IState) => s[prop], [prop]);
+        const mappedState = useMappedState(mapState);
+        if (state[prop] !== mappedState) {
+          flag = true;
+        }
+        return null;
+      };
+
+      render(<Component prop="foo" />);
+      render(<Component prop="bar" />);
+
+      expect(flag).toBe(false);
+    });
+
+    it("doesn't call mapState too often", () => {
+      let mapStateCalls = 0;
+
+      const Component = () => {
+        const mapState = React.useCallback((s: IState) => {
+          mapStateCalls++;
+          return s.foo;
+        }, []);
+        const foo = useMappedState(mapState);
+        return <div>{foo}</div>;
+      };
+
+      render(<Component />);
+      render(<Component />);
+      render(<Component />);
+      render(<Component />);
+
+      // mapState called twice on subscription:
+      // 1. Pull data on initial render
+      // 2. Pull data before right before subscribing in useEffect
+      expect(mapStateCalls).toBe(1);
+    });
   });
 
   describe('useDispatch', () => {

src/create.ts

@@ -1,6 +1,13 @@
 // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
-import {createContext, useContext, useEffect, useRef, useState} from 'react';
+import {
+  createContext,
+  useContext,
+  useEffect,
+  useMemo,
+  useReducer,
+  useRef,
+} from 'react';
 import {Action, Dispatch, Store} from 'redux';
 import shallowEqual from './shallowEqual';
 
@@ -13,6 +20,19 @@ class MissingProviderError extends Error {
   }
 }
 
+function memoizeSingleArg<AT, RT>(fn: (arg: AT) => RT): (arg: AT) => RT {
+  let value: RT;
+  let prevArg: AT;
+
+  return (arg: AT) => {
+    if (prevArg !== arg) {
+      prevArg = arg;
+      value = fn(arg);
+    }
+    return value;
+  };
+}
+
 /**
  * To use redux-react-hook with stronger type safety, or to use with multiple
  * stores in the same app, create() your own instance and re-export the returned
@@ -47,28 +67,31 @@ export function create<
     if (!store) {
       throw new MissingProviderError();
     }
-    const runMapState = () => mapState(store.getState());
-
-    const [derivedState, setDerivedState] = useState(runMapState);
-
-    const lastStore = useRef(store);
-    const lastMapState = useRef(mapState);
-
-    const wrappedSetDerivedState = () => {
-      const newDerivedState = runMapState();
-      setDerivedState(lastDerivedState =>
-        shallowEqual(newDerivedState, lastDerivedState)
-          ? lastDerivedState
-          : newDerivedState,
-      );
-    };
-
-    // If the store or mapState change, rerun mapState
-    if (lastStore.current !== store || lastMapState.current !== mapState) {
-      lastStore.current = store;
-      lastMapState.current = mapState;
-      wrappedSetDerivedState();
-    }
+
+    // We don't keep the derived state but call mapState on every render with current state.
+    // This approach guarantees that useMappedState returns up-to-date derived state.
+    // Since mapState can be expensive and must be a pure function of state we memoize it.
+    const memoizedMapState = useMemo(() => memoizeSingleArg(mapState), [
+      mapState,
+    ]);
+
+    const state = store.getState();
+    const derivedState = memoizedMapState(state);
+
+    // Since we don't keep the derived state we still need to trigger
+    // an update when derived state changes.
+    const [, forceUpdate] = useReducer(x => x + 1, 0);
+
+    // Keep previously commited derived state in a ref.
+    // Compare it to the a one when an action is dispatched
+    // and call forceUpdate if they are different.
+    // We could rely on React's bail out it makes a component
+    // render which is not necessary in that case.
+    const lastStateRef = useRef(derivedState);
+
+    useEffect(() => {
+      lastStateRef.current = derivedState;
+    });
 
     useEffect(() => {
       let didUnsubscribe = false;
@@ -82,7 +105,12 @@ export function create<
           return;
         }
 
-        wrappedSetDerivedState();
+        const newDerivedState = memoizedMapState(store.getState());
+
+        if (!shallowEqual(newDerivedState, lastStateRef.current)) {
+          // In TS definitions userReducer's dispatch requires an argument
+          (forceUpdate as () => void)();
+        }
       };
 
       // Pull data from the store after first render in case the store has
@@ -98,7 +126,7 @@ export function create<
         didUnsubscribe = true;
         unsubscribe();
       };
-    }, [store, mapState]);
+    }, [store, memoizedMapState]);
 
     return derivedState;
   }