Adds memoize option, defaulting to no memoization

src/components/connect.js

@@ -30,16 +30,22 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
   const finalMergeProps = mergeProps || defaultMergeProps
   const shouldUpdateStateProps = finalMapStateToProps.length > 1
   const shouldUpdateDispatchProps = finalMapDispatchToProps.length > 1
-  const { pure = true, withRef = false } = options
+  const { pure = true, withRef = false, memoize = false } = options
+
+  invariant(
+    (memoize === false || typeof memoize === 'function'),
+    '`memoize` must be a function. Instead received %s.',
+    memoize
+  )
 
   // Helps track hot reloading.
   const version = nextVersion++
 
-  function computeStateProps(store, props) {
+  function computeStateProps(mapState, store, props) {
     const state = store.getState()
     const stateProps = shouldUpdateStateProps ?
-      finalMapStateToProps(state, props) :
-      finalMapStateToProps(state)
+      mapState(state, props) :
+      mapState(state)
 
     invariant(
       isPlainObject(stateProps),
@@ -49,11 +55,11 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
     return stateProps
   }
 
-  function computeDispatchProps(store, props) {
+  function computeDispatchProps(mapDispatch, store, props) {
     const { dispatch } = store
     const dispatchProps = shouldUpdateDispatchProps ?
-      finalMapDispatchToProps(dispatch, props) :
-      finalMapDispatchToProps(dispatch)
+      mapDispatch(dispatch, props) :
+      mapDispatch(dispatch)
 
     invariant(
       isPlainObject(dispatchProps),
@@ -115,9 +121,17 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
           `Either wrap the root component in a <Provider>, ` +
           `or explicitly pass "store" as a prop to "${this.constructor.displayName}".`
         )
+      }
 
-        this.stateProps = computeStateProps(this.store, props)
-        this.dispatchProps = computeDispatchProps(this.store, props)
+      componentWillMount() {
+        this.finalMapStateToProps = memoize ?
+          memoize(finalMapStateToProps) :
+          finalMapStateToProps
+        this.finalMapDispatchToProps = memoize ?
+          memoize(finalMapDispatchToProps) :
+          finalMapDispatchToProps
+        this.stateProps = computeStateProps(this.finalMapStateToProps, this.store, this.props)
+        this.dispatchProps = computeDispatchProps(this.finalMapDispatchToProps, this.store, this.props)
         this.state = { storeState: null }
         this.updateState()
       }
@@ -131,7 +145,7 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
       }
 
       updateStateProps(props = this.props) {
-        const nextStateProps = computeStateProps(this.store, props)
+        const nextStateProps = computeStateProps(this.finalMapStateToProps, this.store, props)
         if (shallowEqual(nextStateProps, this.stateProps)) {
           return false
         }
@@ -141,7 +155,7 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
       }
 
       updateDispatchProps(props = this.props) {
-        const nextDispatchProps = computeDispatchProps(this.store, props)
+        const nextDispatchProps = computeDispatchProps(this.finalMapDispatchToProps, this.store, props)
         if (shallowEqual(nextDispatchProps, this.dispatchProps)) {
           return false
         }
@@ -226,6 +240,7 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
         this.version = version
 
         // Update the state and bindings.
+        this.componentWillMount()
         this.trySubscribe()
         this.updateStateProps()
         this.updateDispatchProps()

test/components/connect.spec.js

@@ -4,6 +4,7 @@ import ReactDOM from 'react-dom'
 import TestUtils from 'react-addons-test-utils'
 import { createStore } from 'redux'
 import { connect } from '../../src/index'
+import shallowEqual from '../../src/utils/shallowEqual'
 
 describe('React', () => {
   describe('connect', () => {
@@ -1369,5 +1370,59 @@ describe('React', () => {
       // But render is not because it did not make any actual changes
       expect(renderCalls).toBe(1)
     })
+
+    it('defines finalMapStateToProps on the component', () => {
+
+      const store = createStore(() => ({
+        prefix: 'name: '
+      }))
+      let memoizeHits = 0
+      let memoizeMisses = 0
+      let renderCalls = 0
+
+      function memoize(func) {
+        let lastResult, lastArgs = lastResult = null
+        return (...args) => {
+          if (lastArgs && args.every((value, index) => shallowEqual(value, lastArgs[index]))) {
+            memoizeHits++
+            return lastResult
+          } else if (lastArgs) {
+            memoizeMisses++
+          }
+          lastArgs = args
+          lastResult = func(...args)
+          return lastResult
+        }
+      }
+
+      function selector(state, props) {
+        return { value: props.prefix + state.name }
+      }
+
+      @connect(selector, null, null, { memoize })
+      class Container extends Component {
+        componentDidMount() {
+          this.forceUpdate()
+        }
+        render() {
+          renderCalls++
+          return <div>{this.props.value}</div>
+        }
+      }
+
+      TestUtils.renderIntoDocument(
+        <ProviderMock store={store}>
+          <div>
+            <Container name="one" />
+            <Container name="two" />
+          </div>
+        </ProviderMock>
+      )
+
+      expect(renderCalls).toEqual(4)
+      expect(memoizeHits).toEqual(2)
+      expect(memoizeMisses).toEqual(0)
+    })
+
   })
 })