Fix inconsistent behavior between setState in init and directly assig…

…ning to state (#232)

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Roact Changelog
 
 ## Unreleased Changes
+* Fixed a bug where derived state was lost when assigning directly to state in init ([#232](https://github.com/Roblox/roact/pull/232/))
 * Improved the error message when an invalid changed hook name is used. ([#216](https://github.com/Roblox/roact/pull/216))
 * Fixed a bug where fragments could not be used as children of an element or another fragment. ([#214](https://github.com/Roblox/roact/pull/214))
 

src/Component.lua

@@ -280,6 +280,7 @@ function Component:__mount(reconciler, virtualNode)
 
 	if instance.init ~= nil then
 		instance:init(instance.props)
+		assign(instance.state, instance:__getDerivedState(instance.props, instance.state))
 	end
 
 	-- It's possible for init() to redefine _context!

src/Component.spec/getDerivedStateFromProps.spec.lua

@@ -91,7 +91,12 @@ return function()
 			someState = 2,
 		})
 
-		expect(getDerivedSpy.callCount).to.equal(3)
+		-- getDerivedStateFromProps will be called:
+		-- * Once on empty props
+		-- * Once during the self:setState in init
+		-- * Once more, defensively, on the resulting state AFTER init
+		-- * On updating with new state via updateVirtualNode
+		expect(getDerivedSpy.callCount).to.equal(4)
 
 		local values = getDerivedSpy:captureValues("props", "state")
 
@@ -123,7 +128,11 @@ return function()
 
 		noopReconciler.mountVirtualNode(element, hostParent, hostKey)
 
-		expect(getDerivedSpy.callCount).to.equal(2)
+		-- getDerivedStateFromProps will be called:
+		-- * Once on empty props
+		-- * Once during the self:setState in init
+		-- * Once more, defensively, on the resulting state AFTER init
+		expect(getDerivedSpy.callCount).to.equal(3)
 
 		local values = getDerivedSpy:captureValues("props", "state")
 
@@ -228,4 +237,43 @@ return function()
 		-- getDerivedStateFromProps is always called on initial state
 		expect(stateDerivedSpy.callCount).to.equal(3)
 	end)
+
+	it("should have derived state after assigning to state in init", function()
+		local getStateCallback
+		local getDerivedSpy = createSpy(function()
+			return {
+				derived = true,
+			}
+		end)
+		local WithDerivedState = Component:extend("WithDerivedState")
+
+		WithDerivedState.getDerivedStateFromProps = getDerivedSpy.value
+
+		function WithDerivedState:init()
+			self.state = {
+				init = true,
+			}
+
+			getStateCallback = function()
+				return self.state
+			end
+		end
+
+		function WithDerivedState:render()
+			return nil
+		end
+
+		local hostParent = nil
+		local hostKey = "WithDerivedState"
+		local element = createElement(WithDerivedState)
+
+		noopReconciler.mountVirtualNode(element, hostParent, hostKey)
+
+		expect(getDerivedSpy.callCount).to.equal(2)
+
+		assertDeepEqual(getStateCallback(), {
+			init = true,
+			derived = true,
+		})
+	end)
 end