Overloading connect with factory

src/components/connect.js

@@ -18,51 +18,33 @@ function getDisplayName(WrappedComponent) {
   return WrappedComponent.displayName || WrappedComponent.name || 'Component'
 }
 
+function checkStateShape(stateProps, dispatch) {
+  invariant(
+    isPlainObject(stateProps),
+    '`%sToProps` must return an object. Instead received %s.',
+    dispatch ? 'mapDispatch' : 'mapState',
+    stateProps
+  )
+  return stateProps
+}
+
 // Helps track hot reloading.
 let nextVersion = 0
 
 export default function connect(mapStateToProps, mapDispatchToProps, mergeProps, options = {}) {
   const shouldSubscribe = Boolean(mapStateToProps)
-  const finalMapStateToProps = mapStateToProps || defaultMapStateToProps
-  const finalMapDispatchToProps = isPlainObject(mapDispatchToProps) ?
+  const mapState = mapStateToProps || defaultMapStateToProps
+  const mapDispatch = isPlainObject(mapDispatchToProps) ?
     wrapActionCreators(mapDispatchToProps) :
     mapDispatchToProps || defaultMapDispatchToProps
+
   const finalMergeProps = mergeProps || defaultMergeProps
-  const doStatePropsDependOnOwnProps = finalMapStateToProps.length !== 1
-  const doDispatchPropsDependOnOwnProps = finalMapDispatchToProps.length !== 1
+  const checkMergedEquals = finalMergeProps !== defaultMergeProps
   const { pure = true, withRef = false } = options
 
   // Helps track hot reloading.
   const version = nextVersion++
 
-  function computeStateProps(store, props) {
-    const state = store.getState()
-    const stateProps = doStatePropsDependOnOwnProps ?
-      finalMapStateToProps(state, props) :
-      finalMapStateToProps(state)
-
-    invariant(
-      isPlainObject(stateProps),
-      '`mapStateToProps` must return an object. Instead received %s.',
-      stateProps
-    )
-    return stateProps
-  }
-
-  function computeDispatchProps(store, props) {
-    const { dispatch } = store
-    const dispatchProps = doDispatchPropsDependOnOwnProps ?
-      finalMapDispatchToProps(dispatch, props) :
-      finalMapDispatchToProps(dispatch)
-
-    invariant(
-      isPlainObject(dispatchProps),
-      '`mapDispatchToProps` must return an object. Instead received %s.',
-      dispatchProps
-    )
-    return dispatchProps
-  }
-
   function computeMergedProps(stateProps, dispatchProps, parentProps) {
     const mergedProps = finalMergeProps(stateProps, dispatchProps, parentProps)
     invariant(
@@ -96,8 +78,58 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
         this.clearCache()
       }
 
+      computeStateProps(store, props) {
+        if (!this.finalMapStateToProps) {
+          return this.configureFinalMapState(store, props)
+        }
+
+        const state = store.getState()
+        const stateProps = this.doStatePropsDependOnOwnProps ?
+          this.finalMapStateToProps(state, props) :
+          this.finalMapStateToProps(state)
+
+        return checkStateShape(stateProps)
+      }
+
+      configureFinalMapState(store, props) {
+        const mappedState = mapState(store.getState(), props)
+        const isFactory = typeof mappedState === 'function'
+
+        this.finalMapStateToProps = isFactory ? mappedState : mapState
+        this.doStatePropsDependOnOwnProps = this.finalMapStateToProps.length !== 1
+
+        return isFactory ?
+          this.computeStateProps(store, props) :
+          checkStateShape(mappedState)
+      }
+
+      computeDispatchProps(store, props) {
+        if (!this.finalMapDispatchToProps) {
+          return this.configureFinalMapDispatch(store, props)
+        }
+
+        const { dispatch } = store
+        const dispatchProps = this.doDispatchPropsDependOnOwnProps ?
+          this.finalMapDispatchToProps(dispatch, props) :
+          this.finalMapDispatchToProps(dispatch)
+
+        return checkStateShape(dispatchProps, true)
+      }
+
+      configureFinalMapDispatch(store, props) {
+        const mappedDispatch = mapDispatch(store.dispatch, props)
+        const isFactory = typeof mappedDispatch === 'function'
+
+        this.finalMapDispatchToProps = isFactory ? mappedDispatch : mapDispatch
+        this.doDispatchPropsDependOnOwnProps = this.finalMapDispatchToProps.length !== 1
+
+        return isFactory ?
+          this.computeDispatchProps(store, props) :
+          checkStateShape(mappedDispatch, true)
+      }
+
       updateStatePropsIfNeeded() {
-        const nextStateProps = computeStateProps(this.store, this.props)
+        const nextStateProps = this.computeStateProps(this.store, this.props)
         if (this.stateProps && shallowEqual(nextStateProps, this.stateProps)) {
           return false
         }
@@ -107,7 +139,7 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
       }
 
       updateDispatchPropsIfNeeded() {
-        const nextDispatchProps = computeDispatchProps(this.store, this.props)
+        const nextDispatchProps = this.computeDispatchProps(this.store, this.props)
         if (this.dispatchProps && shallowEqual(nextDispatchProps, this.dispatchProps)) {
           return false
         }
@@ -116,12 +148,14 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
         return true
       }
 
-      updateMergedProps() {
-        this.mergedProps = computeMergedProps(
-          this.stateProps,
-          this.dispatchProps,
-          this.props
-        )
+      updateMergedPropsIfNeeded() {
+        const nextMergedProps = computeMergedProps(this.stateProps, this.dispatchProps, this.props)
+        if (this.mergedProps && checkMergedEquals && shallowEqual(nextMergedProps, this.mergedProps)) {
+          return false
+        }
+
+        this.mergedProps = nextMergedProps
+        return true
       }
 
       isSubscribed() {
@@ -164,6 +198,8 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
         this.haveOwnPropsChanged = true
         this.hasStoreStateChanged = true
         this.renderedElement = null
+        this.finalMapDispatchToProps = null
+        this.finalMapStateToProps = null
       }
 
       handleChange() {
@@ -203,10 +239,10 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
         let shouldUpdateDispatchProps = true
         if (pure && renderedElement) {
           shouldUpdateStateProps = hasStoreStateChanged || (
-            haveOwnPropsChanged && doStatePropsDependOnOwnProps
+            haveOwnPropsChanged && this.doStatePropsDependOnOwnProps
           )
           shouldUpdateDispatchProps =
-            haveOwnPropsChanged && doDispatchPropsDependOnOwnProps
+            haveOwnPropsChanged && this.doDispatchPropsDependOnOwnProps
         }
 
         let haveStatePropsChanged = false
@@ -224,7 +260,7 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
           haveDispatchPropsChanged ||
           haveOwnPropsChanged
         ) {
-          this.updateMergedProps()
+          haveMergedPropsChanged = this.updateMergedPropsIfNeeded()
         } else {
           haveMergedPropsChanged = false
         }

test/components/connect.spec.js

@@ -39,7 +39,7 @@ describe('React', () => {
       @connect()
       class Container extends Component {
         render() {
-          return <div {...this.props} />
+          return <Passthrough {...this.props} />
         }
       }
 
@@ -1509,5 +1509,135 @@ describe('React', () => {
       // But render is not because it did not make any actual changes
       expect(renderCalls).toBe(1)
     })
+
+    it('should allow providing a factory function to mapStateToProps', () => {
+      let updatedCount = 0
+      let memoizedReturnCount = 0
+      const store = createStore(() => ({ value: 1 }))
+
+      const mapStateFactory = () => {
+        let lastProp, lastVal, lastResult
+        return (state, props) => {
+          if (props.name === lastProp && lastVal === state.value) {
+            memoizedReturnCount++
+            return lastResult
+          }
+          lastProp = props.name
+          lastVal = state.value
+          return lastResult = { someObject: { prop: props.name, stateVal: state.value } }
+        }
+      }
+
+      @connect(mapStateFactory)
+      class Container extends Component {
+        componentWillUpdate() {
+          updatedCount++
+        }
+        render() {
+          return <Passthrough {...this.props} />
+        }
+      }
+
+      TestUtils.renderIntoDocument(
+        <ProviderMock store={store}>
+          <div>
+            <Container name="a" />
+            <Container name="b" />
+          </div>
+        </ProviderMock>
+      )
+
+      store.dispatch({ type: 'test' })
+      expect(updatedCount).toBe(0)
+      expect(memoizedReturnCount).toBe(2)
+    })
+
+    it('should allow providing a factory function to mapDispatchToProps', () => {
+      let updatedCount = 0
+      let memoizedReturnCount = 0
+      const store = createStore(() => ({ value: 1 }))
+
+      const mapDispatchFactory = () => {
+        let lastProp, lastResult
+        return (dispatch, props) => {
+          if (props.name === lastProp) {
+            memoizedReturnCount++
+            return lastResult
+          }
+          lastProp = props.name
+          return lastResult = { someObject: { dispatchFn: dispatch } }
+        }
+      }
+      function mergeParentDispatch(stateProps, dispatchProps, parentProps) {
+        return { ...stateProps, ...dispatchProps, name: parentProps.name }
+      }
+
+      @connect(null, mapDispatchFactory, mergeParentDispatch)
+      class Passthrough extends Component {
+        componentWillUpdate() {
+          updatedCount++
+        }
+        render() {
+          return <div {...this.props} />
+        }
+      }
+
+      class Container extends Component {
+        constructor(props) {
+          super(props)
+          this.state = { count: 0 }
+        }
+        componentDidMount() {
+          this.setState({ count: 1 })
+        }
+        render() {
+          const { count } = this.state
+          return (
+            <div>
+              <Passthrough count={count} name="a" />
+              <Passthrough count={count} name="b" />
+            </div>
+          )
+        }
+      }
+
+      TestUtils.renderIntoDocument(
+        <ProviderMock store={store}>
+          <Container />
+        </ProviderMock>
+      )
+
+      store.dispatch({ type: 'test' })
+      expect(updatedCount).toBe(0)
+      expect(memoizedReturnCount).toBe(2)
+    })
+
+    it('should not call update if mergeProps return value has not changed', () => {
+      let mapStateCalls = 0
+      let renderCalls = 0
+      const store = createStore(stringBuilder)
+
+      @connect(() => ({ a: ++mapStateCalls }), null, () => ({ changed: false }))
+      class Container extends Component {
+        render() {
+          renderCalls++
+          return <Passthrough {...this.props} />
+        }
+      }
+
+      TestUtils.renderIntoDocument(
+        <ProviderMock store={store}>
+          <Container />
+        </ProviderMock>
+      )
+
+      expect(renderCalls).toBe(1)
+      expect(mapStateCalls).toBe(1)
+
+      store.dispatch({ type: 'APPEND', body: 'a' })
+
+      expect(mapStateCalls).toBe(2)
+      expect(renderCalls).toBe(1)
+    })
   })
 })