Allow functions to be passed to setState

lib/Component.lua

@@ -121,6 +121,11 @@ function Component:setState(partialState)
 		error(INVALID_SETSTATE_MESSAGE, 0)
 	end
 
+	-- If the partial state is a function, invoke it to get the actual partial state.
+	if type(partialState) == "function" then
+		partialState = partialState(self.state, self.props)
+	end
+
 	local newState = {}
 
 	for key, value in pairs(self.state) do

lib/Component.spec.lua

@@ -326,5 +326,50 @@ return function()
 
 			Reconciler.teardown(instance)
 		end)
+
+		it("should invoke functions to compute a partial state", function()
+			local TestComponent = Component:extend("TestComponent")
+			local setStateCallback, getStateCallback, getPropsCallback
+
+			function TestComponent:init()
+				setStateCallback = function(newState)
+					self:setState(newState)
+				end
+
+				getStateCallback = function()
+					return self.state
+				end
+
+				getPropsCallback = function()
+					return self.props
+				end
+
+				self.state = {
+					value = 0
+				}
+			end
+
+			function TestComponent:render()
+				return nil
+			end
+
+			local element = Core.createElement(TestComponent)
+			local instance = Reconciler.reify(element)
+
+			expect(getStateCallback().value).to.equal(0)
+
+			setStateCallback(function(state, props)
+				expect(state).to.equal(getStateCallback())
+				expect(props).to.equal(getPropsCallback())
+
+				return {
+					value = state.value + 1
+				}
+			end)
+
+			expect(getStateCallback().value).to.equal(1)
+
+			Reconciler.teardown(instance)
+		end)
 	end)
 end