From 8f5f825e17602c42b4cf4897ff4dda3246823ffb Mon Sep 17 00:00:00 2001
From: Esa-Matti Suuronen <esa-matti@suuronen.org>
Date: Wed, 9 Sep 2015 10:34:39 +0300
Subject: [PATCH] Memoize mapState

mapState can be memoized automatically when it does not take the second
parameter.

I had to change invocationCount assert to 1 from 2. I'm pretty sure it
should be invoked only once on the initial render and not on prop
changes again.
---
 src/components/createConnect.js | 19 +++++++++++++++++--
 test/components/connect.spec.js |  2 +-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/components/createConnect.js b/src/components/createConnect.js
index d4e7f3c28..163b5a10c 100644
--- a/src/components/createConnect.js
+++ b/src/components/createConnect.js
@@ -108,6 +108,7 @@ export default function createConnect(React) {
           super(props, context);
           this.version = version;
           this.store = props.store || context.store;
+          this.mapStateMemoize = null;
 
           invariant(this.store,
             `Could not find "store" in either the context or ` +
@@ -157,6 +158,9 @@ export default function createConnect(React) {
               this.updateDispatchProps(nextProps);
             }
 
+            if (shouldUpdateStateProps) {
+              this.mapStateMemoize = null;
+            }
           }
         }
 
@@ -169,7 +173,12 @@ export default function createConnect(React) {
             return;
           }
 
-          this.setState({storeState: this.store.getState()});
+          const storeState = this.store.getState();
+          if (storeState !== this.state.storeState) {
+            this.mapStateMemoize = null;
+          }
+
+          this.setState({storeState});
         }
 
         getWrappedInstance() {
@@ -181,8 +190,12 @@ export default function createConnect(React) {
         }
 
         computeNextState() {
+          if (this.mapStateMemoize === null) {
+            this.mapStateMemoize = computeStateProps(this.state.storeState, this.props);
+          }
+
           return computeNextState(
-            computeStateProps(this.state.storeState, this.props),
+            this.mapStateMemoize,
             this.state.dispatchProps,
             this.props
           );
@@ -214,6 +227,8 @@ export default function createConnect(React) {
           // We are hot reloading!
           this.version = version;
 
+          this.mapStateMemoize = null;
+
           // Update the state and bindings.
           this.trySubscribe();
           this.updateDispatchProps();
diff --git a/test/components/connect.spec.js b/test/components/connect.spec.js
index 367228bd4..a5c1fdece 100644
--- a/test/components/connect.spec.js
+++ b/test/components/connect.spec.js
@@ -441,7 +441,7 @@ describe('React', () => {
       outerComponent.setFoo('BAR');
       outerComponent.setFoo('DID');
 
-      expect(invocationCount).toEqual(2);
+      expect(invocationCount).toEqual(1);
     });
 
     it('should invoke mapState every time props are changed if it has a second argument', () => {