Allow setState inside willUpdate and init

CHANGELOG.md

@@ -11,6 +11,7 @@
 * Added `createRef` ([#70](https://github.com/Roblox/roact/issues/70), [#92](https://github.com/Roblox/roact/pull/92))
 * Added a warning when an element changes type during reconciliation ([#88](https://github.com/Roblox/roact/issues/88), [#137](https://github.com/Roblox/roact/pull/137))
 * Ref switching now occurs in one pass, which should fix edge cases where the result of a ref is `nil`, especially in property changed events ([#98](https://github.com/Roblox/roact/pull/98))
+* `setState` can now be called inside `init` and `willUpdate`. Instead of triggering a new render, it will affect the currently scheduled one. ([#139](https://github.com/Roblox/roact/pull/139))
 
 ## 1.0.0 Prerelease 2 (March 22, 2018)
 * Removed `is*Element` methods, this is unlikely to affect anyone ([#50](https://github.com/Roblox/roact/pull/50))

docs/api-reference.md

@@ -203,7 +203,18 @@ init(initialProps) -> void
 
 `init` is called exactly once when a new instance of a component is created. It can be used to set up the initial `state`, as well as any non-`render` related values directly on the component.
 
-`init` is the only place where you can assign to `state` directly, as opposed to using `setState`:
+Use `setState` inside of `init` to set up your initial component state:
+
+```lua
+function MyComponent:init()
+	self:setState({
+		position = 0,
+		velocity = 10
+	})
+end
+```
+
+In older versions of Roact, `setState` was disallowed in `init`, and you would instead assign to `state` directly. It's simpler to use `setState`, but assigning directly to `state` is still acceptable inside `init`:
 
 ```lua
 function MyComponent:init()
@@ -281,6 +292,16 @@ end
 
 Setting a field in the state to `Roact.None` will clear it from the state. This is the only way to remove a field from a component's state!
 
+!!! warning
+	`setState` can be called from anywhere **except**:
+
+	* Lifecycle hooks: `willUnmount`
+	* Pure functions: `render`, `shouldUpdate`
+
+	Calling `setState` inside of `init` or `willUpdate` has special behavior. Because Roact is already going to update a component in these cases, that update will be replaced instead of another being scheduled.
+
+	Roact may support calling `setState` in currently-disallowed places in the future.
+
 !!! warning
 	**`setState` does not always resolve synchronously!** Roact may batch and reschedule state updates in order to reduce the number of total renders.
 
@@ -291,13 +312,6 @@ Setting a field in the state to `Roact.None` will clear it from the state. This
 	* [RFClarification: why is `setState` asynchronous?](https://github.com/facebook/react/issues/11527#issuecomment-360199710)
 	* [Does React keep the order for state updates?](https://stackoverflow.com/a/48610973/802794)
 
-!!! warning
-	Calling `setState` from any of these places is not allowed at this time and will throw an error:
-
-	* Lifecycle hooks: `willUpdate`, `willUnmount`
-	* Initialization: `init`
-	* Pure functions: `render`, `shouldUpdate`
-
 ### shouldUpdate
 ```
 shouldUpdate(nextProps, nextState) -> bool
@@ -350,6 +364,8 @@ willUpdate(nextProps, nextState) -> void
 
 `willUpdate` is fired after an update is started but before a component's state and props are updated.
 
+`willUpdate` can be used to make tweaks to your component's state using `setState`. Often, this should be done in `getDerivedStateFromProps` instead.
+
 ### didUpdate
 ```
 didUpdate(previousProps, previousState) -> void

docs/guide/state-and-lifecycle.md

@@ -3,15 +3,48 @@ In the previous section, we talked about using components to create reusable chu
 Stateful components do everything that functional components do, but have the addition of mutable *state* and *lifecycle methods*.
 
 ## State
+**State** is the term we use to talk about values that are owned by a component itself.
 
-!!! info
-	This section is incomplete!
+Unlike **props**, which are passed to a component from above, **state** is created within a component and can only be updated by that component.
+
+We can set up the initial state of a stateful component inside of a method named `init`:
+
+```lua
+function MyComponent:init()
+	self:setState({
+		currentTime = 0
+	})
+end
+```
+
+To update state, we use a special method named `setState`. `setState` will merge any values we give it into our state. It will overwrite any existing values, and leave any values we don't specify alone.
+
+There's another form of `setState` we can use. When the new state we want our component to have depends on our current state, like incrementing a value, we use this form:
+
+```lua
+-- This is another special method, didMount, that we'll talk about in a moment.
+function MyComponent:didMount()
+	self:setState(function(state)
+		return {
+			currentTime = currentTime + state.currentTime
+		}
+	end)
+end
+```
+
+In this case, we're passing a _function_ to `setState`. This function is called and passed the current state, and returns a new state. It can also return `nil` to abort the state update, which lets Roact make some handy optimizations.
+
+Right now, this version of `setState` works exactly the same way as the version that accepts an object. In the future, Roact will support optimizations that make this difference more important, like [asynchronous rendering](https://github.com/Roblox/roact/issues/18).
 
 ## Lifecycle Methods
 Stateful components can provide methods to Roact that are called when certain things happen to a component instance.
 
 Lifecycle methods are a great place to send off network requests, measure UI ([with the help of refs](/advanced/refs)), wrap non-Roact components, and produce other side-effects.
 
+The most useful lifecycle methods are generally `didMount` and `didUpdate`. Most components that do things that are difficult to express in Roact itself will use these lifecycle methods.
+
+Here's a chart of all of the methods available. You can also check out the [Lifecycle Methods](../api-reference/#lifecycle-methods) section of the API reference for more details.
+
 <div align="center">
 	<a href="../../images/lifecycle.svg">
 		<img src="../../images/lifecycle.svg" alt="Diagram of Roact Lifecycle" />
@@ -32,11 +65,10 @@ local Roact = require(ReplicatedStorage.Roact)
 local Clock = Roact.Component:extend("Clock")
 
 function Clock:init()
-	-- In init, you should assign to 'state' directly.
-	-- Use this opportunity to set any initial values.
-	self.state = {
+	-- In init, we can use setState to set up our initial component state.
+	self:setState({
 		currentTime = 0
-	}
+	})
 end
 
 -- This render function is almost completely unchanged from the first example.

lib/Component.lua

@@ -92,6 +92,11 @@ function Component:extend(name)
 		-- You can see a list of reasons in invalidSetStateMessages.
 		self._setStateBlockedReason = nil
 
+		-- When set to true, setState should not trigger an update, but should
+		-- instead just update self.state. Lifecycle events like `willUpdate`
+		-- can set this to change the behavior of setState slightly.
+		self._setStateWithoutUpdate = false
+
 		if class.defaultProps == nil then
 			self.props = passedProps
 		else
@@ -109,16 +114,13 @@ function Component:extend(name)
 
 		setmetatable(self, class)
 
+		self.state = {}
+
 		-- Call the user-provided initializer, where state and _props are set.
 		if class.init then
-			self._setStateBlockedReason = "init"
+			self._setStateWithoutUpdate = true
 			class.init(self, self.props)
-			self._setStateBlockedReason = nil
-		end
-
-		-- The user constructer might not set state, so we can.
-		if not self.state then
-			self.state = {}
+			self._setStateWithoutUpdate = false
 		end
 
 		if class.getDerivedStateFromProps then
@@ -236,7 +238,12 @@ function Component:setState(partialState)
 	end
 
 	local newState = merge(self.state, partialState)
-	self:_update(nil, newState)
+
+	if self._setStateWithoutUpdate then
+		self.state = newState
+	else
+		self:_update(nil, newState)
+	end
 end
 
 --[[
@@ -314,9 +321,9 @@ end
 ]]
 function Component:_forceUpdate(newProps, newState)
 	if self.willUpdate then
-		self._setStateBlockedReason = "willUpdate"
+		self._setStateWithoutUpdate = true
 		self:willUpdate(newProps or self.props, newState or self.state)
-		self._setStateBlockedReason = nil
+		self._setStateWithoutUpdate = false
 	end
 
 	local oldProps = self.props

lib/Component.spec.lua

@@ -366,26 +366,6 @@ return function()
 	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
-
-			function InitComponent:render()
-				return nil
-			end
-
-			local initElement = createElement(InitComponent)
-
-			expect(function()
-				Reconciler.mount(initElement)
-			end).to.throw()
-		end)
-
 		it("should throw when called in render", function()
 			local RenderComponent = Component:extend("RenderComponent")
 
@@ -433,7 +413,28 @@ return function()
 			end).to.throw()
 		end)
 
-		it("should throw when called in willUpdate", function()
+		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 = createElement(TestComponent)
+			local instance = Reconciler.mount(element)
+
+			expect(function()
+				Reconciler.unmount(instance)
+			end).to.throw()
+		end)
+
+		it("should only render once when called in willUpdate", function()
 			local TestComponent = Component:extend("TestComponent")
 			local forceUpdate
 
@@ -443,7 +444,9 @@ return function()
 				end
 			end
 
+			local renderCount = 0
 			function TestComponent:render()
+				renderCount = renderCount + 1
 				return nil
 			end
 
@@ -455,31 +458,39 @@ return function()
 
 			local testElement = createElement(TestComponent)
 
-			expect(function()
-				Reconciler.mount(testElement)
-				forceUpdate()
-			end).to.throw()
+			local handle = Reconciler.mount(testElement)
+
+			expect(renderCount).to.equal(1)
+
+			forceUpdate()
+
+			expect(renderCount).to.equal(2)
+
+			Reconciler.unmount(handle)
 		end)
 
-		it("should throw when called in willUnmount", function()
+		it("should only render once when called in init", function()
 			local TestComponent = Component:extend("TestComponent")
 
+			function TestComponent:init()
+				self:setState({
+					a = 7,
+				})
+			end
+
+			local renderCount = 0
 			function TestComponent:render()
+				renderCount = renderCount + 1
 				return nil
 			end
 
-			function TestComponent:willUnmount()
-				self:setState({
-					a = 1
-				})
-			end
+			local testElement = createElement(TestComponent)
 
-			local element = createElement(TestComponent)
-			local instance = Reconciler.mount(element)
+			local handle = Reconciler.mount(testElement)
 
-			expect(function()
-				Reconciler.unmount(instance)
-			end).to.throw()
+			expect(renderCount).to.equal(1)
+
+			Reconciler.unmount(handle)
 		end)
 
 		it("should remove values from state when the value is Core.None", function()