setState sanity checking

lib/Component.lua

@@ -43,6 +43,8 @@ function Component:extend(name)
 		local self = {}
 
 		self.props = props
+		-- Used for tracking whether the component is in a position to set state.
+		self._canSetState = false
 		self._context = {}
 
 		-- Shallow copy all context values from our parent element.
@@ -64,6 +66,9 @@ function Component:extend(name)
 			self.state = {}
 		end
 
+		-- Now that state has definitely been set, we can now allow it to be changed.
+		self._canSetState = true
+
 		return self
 	end
 
@@ -101,6 +106,14 @@ end
 	current state object.
 ]]
 function Component:setState(partialState)
+	-- State cannot be set in any of the following places:
+	-- * During the component's init function
+	-- * During the component's render function
+	-- * After the component has been unmounted (or is in the process of unmounting, e.g. willUnmount)
+	if not self._canSetState then
+		error("setState cannot be used currently: are you calling setState from an init, render, or willUnmount function?", 0)
+	end
+
 	local newState = {}
 
 	for key, value in pairs(self.state) do
@@ -149,7 +162,9 @@ function Component:_forceUpdate(newProps, newState)
 		self.state = newState
 	end
 
+	self._canSetState = false
 	local newChildElement = self:render()
+	self._canSetState = true
 
 	if self._handle._reified ~= nil then
 		-- We returned an element before, update it.

lib/Component.spec.lua

@@ -173,4 +173,59 @@ return function()
 
 		Reconciler.teardown(instance)
 	end)
+
+	describe("setState", function()
+		it("should throw when called in init", function()
+			local InitComponent = Component:extend("InitComponent")
+
+			function InitComponent:init()
+				self:setState({
+					a = 1
+				})
+			end
+
+			local initElement = Core.createElement(InitComponent)
+
+			expect(function()
+				Reconciler.reify(initElement)
+			end).to.throw()
+		end)
+
+		it("should throw when called in render", function()
+			local RenderComponent = Component:extend("RenderComponent")
+
+			function RenderComponent:render()
+				self:setState({
+					a = 1
+				})
+			end
+
+			local renderElement = Core.createElement(RenderComponent)
+
+			expect(function()
+				Reconciler.reify(renderElement)
+			end).to.throw()
+		end)
+
+		it("should throw when called in willUnmount", function()
+			local TestComponent = Component:extend("TestComponent")
+
+			function TestComponent:render()
+				return nil
+			end
+
+			function TestComponent:willUnmount()
+				self:setState({
+					a = 1
+				})
+			end
+
+			local element = Core.createElement(TestComponent)
+			local instance = Reconciler.reify(element)
+
+			expect(function()
+				Reconciler.teardown(instance)
+			end).to.throw()
+		end)
+	end)
 end

lib/Reconciler.lua

@@ -60,6 +60,9 @@ function Reconciler.teardown(instanceHandle)
 			Reconciler.teardown(instanceHandle._reified)
 		end
 	elseif Core.isStatefulElement(element) then
+		-- Stop the component from setting state in willUnmount or anywhere thereafter.
+		instanceHandle._instance._canSetState = false
+
 		-- Tell the component we're about to tear everything down.
 		-- This gives it some notice!
 		if instanceHandle._instance.willUnmount then