@@ -30,16 +30,22 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
30
30
const finalMergeProps = mergeProps || defaultMergeProps
31
31
const shouldUpdateStateProps = finalMapStateToProps . length > 1
32
32
const shouldUpdateDispatchProps = finalMapDispatchToProps . length > 1
33
- const { pure = true , withRef = false } = options
33
+ const { pure = true , withRef = false , memoize = false } = options
34
+
35
+ invariant (
36
+ ( memoize === false || typeof memoize === 'function' ) ,
37
+ '`memoize` must be a function. Instead received %s.' ,
38
+ memoize
39
+ )
34
40
35
41
// Helps track hot reloading.
36
42
const version = nextVersion ++
37
43
38
- function computeStateProps ( store , props ) {
44
+ function computeStateProps ( mapState , store , props ) {
39
45
const state = store . getState ( )
40
46
const stateProps = shouldUpdateStateProps ?
41
- finalMapStateToProps ( state , props ) :
42
- finalMapStateToProps ( state )
47
+ mapState ( state , props ) :
48
+ mapState ( state )
43
49
44
50
invariant (
45
51
isPlainObject ( stateProps ) ,
@@ -49,11 +55,11 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
49
55
return stateProps
50
56
}
51
57
52
- function computeDispatchProps ( store , props ) {
58
+ function computeDispatchProps ( mapDispatch , store , props ) {
53
59
const { dispatch } = store
54
60
const dispatchProps = shouldUpdateDispatchProps ?
55
- finalMapDispatchToProps ( dispatch , props ) :
56
- finalMapDispatchToProps ( dispatch )
61
+ mapDispatch ( dispatch , props ) :
62
+ mapDispatch ( dispatch )
57
63
58
64
invariant (
59
65
isPlainObject ( dispatchProps ) ,
@@ -115,9 +121,17 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
115
121
`Either wrap the root component in a <Provider>, ` +
116
122
`or explicitly pass "store" as a prop to "${ this . constructor . displayName } ".`
117
123
)
124
+ }
118
125
119
- this . stateProps = computeStateProps ( this . store , props )
120
- this . dispatchProps = computeDispatchProps ( this . store , props )
126
+ componentWillMount ( ) {
127
+ this . finalMapStateToProps = memoize ?
128
+ memoize ( finalMapStateToProps ) :
129
+ finalMapStateToProps
130
+ this . finalMapDispatchToProps = memoize ?
131
+ memoize ( finalMapDispatchToProps ) :
132
+ finalMapDispatchToProps
133
+ this . stateProps = computeStateProps ( this . finalMapStateToProps , this . store , this . props )
134
+ this . dispatchProps = computeDispatchProps ( this . finalMapDispatchToProps , this . store , this . props )
121
135
this . state = { storeState : null }
122
136
this . updateState ( )
123
137
}
@@ -131,7 +145,7 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
131
145
}
132
146
133
147
updateStateProps ( props = this . props ) {
134
- const nextStateProps = computeStateProps ( this . store , props )
148
+ const nextStateProps = computeStateProps ( this . finalMapStateToProps , this . store , props )
135
149
if ( shallowEqual ( nextStateProps , this . stateProps ) ) {
136
150
return false
137
151
}
@@ -141,7 +155,7 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
141
155
}
142
156
143
157
updateDispatchProps ( props = this . props ) {
144
- const nextDispatchProps = computeDispatchProps ( this . store , props )
158
+ const nextDispatchProps = computeDispatchProps ( this . finalMapDispatchToProps , this . store , props )
145
159
if ( shallowEqual ( nextDispatchProps , this . dispatchProps ) ) {
146
160
return false
147
161
}
@@ -226,6 +240,7 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
226
240
this . version = version
227
241
228
242
// Update the state and bindings.
243
+ this . componentWillMount ( )
229
244
this . trySubscribe ( )
230
245
this . updateStateProps ( )
231
246
this . updateDispatchProps ( )
0 commit comments