Add stateThunk / dispatchThunk options

src/components/connect.js

@@ -22,6 +22,7 @@ function getDisplayName(WrappedComponent) {
 let nextVersion = 0
 
 export default function connect(mapStateToProps, mapDispatchToProps, mergeProps, options = {}) {
+  const { pure = true, withRef = false } = options
   const shouldSubscribe = Boolean(mapStateToProps)
   const finalMapStateToProps = mapStateToProps || defaultMapStateToProps
   const finalMapDispatchToProps = isPlainObject(mapDispatchToProps) ?
@@ -30,16 +31,15 @@ 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
 
   // 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 +49,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),
@@ -116,12 +116,18 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
           `or explicitly pass "store" as a prop to "${this.constructor.displayName}".`
         )
 
-        this.stateProps = computeStateProps(this.store, props)
-        this.dispatchProps = computeDispatchProps(this.store, props)
+        this.assignMapFunctions()
+        this.stateProps = computeStateProps(this.finalMapStateToProps, this.store, props)
+        this.dispatchProps = computeDispatchProps(this.finalMapDispatchToProps, this.store, props)
         this.state = { storeState: null }
         this.updateState()
       }
 
+      assignMapFunctions() {
+        this.finalMapStateToProps = finalMapStateToProps
+        this.finalMapDispatchToProps = finalMapDispatchToProps
+      }
+
       computeNextState(props = this.props) {
         return computeNextState(
           this.stateProps,
@@ -131,7 +137,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 +147,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 +232,7 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
         this.version = version
 
         // Update the state and bindings.
+        this.assignMapFunctions()
         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', () => {
@@ -190,7 +191,7 @@ describe('React', () => {
           return (
             <ProviderMock store={store}>
               <ConnectContainer bar={this.state.bar} />
-             </ProviderMock>
+            </ProviderMock>
           )
         }
       }
@@ -1369,5 +1370,68 @@ 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 memoizer(WrappedComponent) {
+        let assignMapFunctions = WrappedComponent.prototype.assignMapFunctions
+        WrappedComponent.prototype.assignMapFunctions = function () {
+          assignMapFunctions.call(this)
+          this.finalMapStateToProps = memoize(this.finalMapStateToProps)
+        }
+        return WrappedComponent
+      }
+
+      function computeValue(state, props) {
+        return { value: props.prefix + state.name }
+      }
+
+      @memoizer
+      @connect(computeValue, null, null, { stateThunk: true })
+      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)
+    })
   })
 })