Rewrite connect2 to be a lot simpler

lib/connect2.lua

@@ -14,20 +14,46 @@ local function noop()
 	return nil
 end
 
+--[[
+	The stateUpdater accepts props when they update and computes the
+	complete set of props that should be passed to the wrapped component.
+
+	Each connected component will have a stateUpdater created for it.
+
+	stateUpdater is put into the component's state in order for
+	getDerivedStateFromProps to be able to access it. It is not mutated.
+]]
+local function makeStateUpdater(store, mapStateToProps)
+	return function(nextProps, prevState, mappedStoreState)
+		-- The caller can optionally provide mappedStoreState if it needed that
+		-- value beforehand. Doing so is purely an optimization.
+		if mappedStoreState == nil then
+			mappedStoreState = mapStateToProps(store:getState(), nextProps)
+		end
+
+		local propsForChild = join(nextProps, mappedStoreState, prevState.mappedStoreDispatch)
+
+		return {
+			mappedStoreState = mappedStoreState,
+			propsForChild = propsForChild,
+		}
+	end
+end
+
 --[[
 	mapStateToProps:
 		(storeState, props) -> partialProps
 		OR
 		() -> (storeState, props) -> partialProps
 	mapDispatchToProps: (dispatch) -> partialProps
 ]]
-local function connect(mapStateToProps, mapDispatchToProps)
+local function connect(mapStateToPropsOrThunk, mapDispatchToProps)
 	local connectTrace = debug.traceback()
 
-	if mapStateToProps ~= nil then
-		assert(typeof(mapStateToProps) == "function", "mapStateToProps must be a function or nil!")
+	if mapStateToPropsOrThunk ~= nil then
+		assert(typeof(mapStateToPropsOrThunk) == "function", "mapStateToProps must be a function or nil!")
 	else
-		mapStateToProps = noop
+		mapStateToPropsOrThunk = noop
 	end
 
 	if mapDispatchToProps ~= nil then
@@ -51,19 +77,13 @@ local function connect(mapStateToProps, mapDispatchToProps)
 
 		local componentName = ("RoduxConnection(%s)"):format(tostring(innerComponent))
 
-		local outerComponent = Roact.Component:extend(componentName)
+		local Connection = Roact.Component:extend(componentName)
 
-		function outerComponent.getDerivedStateFromProps(nextProps, prevState)
-			local stateValues = prevState.stateMapper(prevState.storeState, nextProps)
-			local combinedValues = join(nextProps, stateValues, prevState.dispatchValues)
-
-			return {
-				stateValues = stateValues,
-				combinedValues = combinedValues,
-			}
+		function Connection.getDerivedStateFromProps(nextProps, prevState)
+			return prevState.stateUpdater(nextProps, prevState)
 		end
 
-		function outerComponent:init()
+		function Connection:init()
 			self.store = self._context[storeKey]
 
 			if self.store == nil then
@@ -80,77 +100,74 @@ local function connect(mapStateToProps, mapDispatchToProps)
 
 			local storeState = self.store:getState()
 
-			local stateMapper = mapStateToProps
-			local stateValues = mapStateToProps(storeState, self.props)
+			local mapStateToProps = mapStateToPropsOrThunk
+			local mappedStoreState = mapStateToProps(storeState, self.props)
 
-			-- mapStateToProps can return a function instead of a state value.
-			-- In this variant, we keep that value as our 'state mapper' instead
-			-- of the original mapStateToProps. This matches react-redux and
-			-- enables connectors to keep instance-level state.
-			if typeof(stateValues) == "function" then
-				stateMapper = stateValues
-				stateValues = stateValues(storeState, self.props)
+			-- mapStateToPropsOrThunk can return a function instead of a state
+			-- value. In this variant, we keep that value as mapStateToProps
+			-- instead of the original mapStateToProps. This matches react-redux
+			-- and enables connectors to keep instance-level state.
+			if typeof(mappedStoreState) == "function" then
+				mapStateToProps = mappedStoreState
+				mappedStoreState = mapStateToProps(storeState, self.props)
 			end
 
-			if stateValues ~= nil and typeof(stateValues) ~= "table" then
+			if mappedStoreState ~= nil and typeof(mappedStoreState) ~= "table" then
 				local message = formatMessage({
 					"mapStateToProps must either return a table, or return another function that returns a table.",
 					"Instead, it returned %q, which is of type %s.",
 				}, {
-					tostring(stateValues),
-					typeof(stateValues),
+					tostring(mappedStoreState),
+					typeof(mappedStoreState),
 				})
 
 				error(message)
 			end
 
-			local dispatchValues = mapDispatchToProps(function(...)
+			local mappedStoreDispatch = mapDispatchToProps(function(...)
 				return self.store:dispatch(...)
 			end)
 
+			local stateUpdater = makeStateUpdater(self.store, mapStateToProps)
+
+			self.mapStateToProps = mapStateToProps
+
+			-- All of the values that we put into state are intended to be used
+			-- by getDerivedStateFromProps, constructed using makeStateUpdater.
 			self.state = {
-				storeState = storeState,
-				stateMapper = stateMapper,
-				stateValues = stateValues,
-				dispatchValues = dispatchValues,
-				combinedValues = join(self.props, stateValues, dispatchValues)
+				stateUpdater = stateUpdater,
+				mappedStoreDispatch = mappedStoreDispatch,
 			}
-		end
 
-		function outerComponent:updateState(newStoreState)
-			self:setState(function(prevState, props)
-				local newStateValues = prevState.stateMapper(newStoreState, props)
-
-				-- We don't need to update if the result was the same.
-				if shallowEqual(newStateValues, prevState.stateValues) then
-					return nil
-				end
+			self.state.propsForChild = stateUpdater(self.props, self.state, mappedStoreState)
+		end
 
-				local newCombinedState = join(props, newStateValues, prevState.dispatchValues)
+		function Connection:didMount()
+			self.storeChangedConnection = self.store.changed:connect(function(storeState)
+				self:setState(function(prevState, props)
+					local mappedStoreState = self.mapStateToProps(storeState, props)
 
-				return {
-					storeState = newStoreState,
-					stateValues = newStateValues,
-					combinedValues = newCombinedState,
-				}
-			end)
-		end
+					-- We run this check here so that we only check shallow
+					-- equality with the result of mapStateToProps, and not the
+					-- other props that could be passed through the connector.
+					if shallowEqual(mappedStoreState, prevState.mappedStoreState) then
+						return nil
+					end
 
-		function outerComponent:didMount()
-			self.eventHandle = self.store.changed:connect(function(storeState)
-				self:updateState(storeState)
+					return prevState.stateUpdater(props, prevState, mappedStoreState)
+				end)
 			end)
 		end
 
-		function outerComponent:willUnmount()
-			self.eventHandle:disconnect()
+		function Connection:willUnmount()
+			self.storeChangedConnection:disconnect()
 		end
 
-		function outerComponent:render()
-			return Roact.createElement(innerComponent, self.state.combinedValues)
+		function Connection:render()
+			return Roact.createElement(innerComponent, self.state.propsForChild)
 		end
 
-		return outerComponent
+		return Connection
 	end
 end