Merge 28cc872 into deaa800

src/Component.lua

@@ -196,6 +196,46 @@ function Component:render()
 	error(message, 0)
 end
 
+--[[
+	Retrieves the context value corresponding to the given key. Can return nil
+	if a requested context key is not present
+]]
+function Component:__getContext(key)
+	if config.internalTypeChecks then
+		internalAssert(Type.of(self) == Type.StatefulComponentInstance, "Invalid use of `__getContext`")
+		internalAssert(key ~= nil, "Context key cannot be nil")
+	end
+
+	local virtualNode = self[InternalData].virtualNode
+	local context = virtualNode.context
+
+	return context[key]
+end
+
+--[[
+	Adds a new context entry to this component's context table (which will be
+	passed down to child components).
+]]
+function Component:__addContext(key, value)
+	if config.internalTypeChecks then
+		internalAssert(Type.of(self) == Type.StatefulComponentInstance, "Invalid use of `__addContext`")
+	end
+	local virtualNode = self[InternalData].virtualNode
+
+	-- Make sure we store a reference to the component's original, unmodified
+	-- context the virtual node. In the reconciler, we'll restore the original
+	-- context if we need to replace the node (this happens when a node gets
+	-- re-rendered as a different component)
+	if virtualNode.originalContext == nil then
+		virtualNode.originalContext = virtualNode.context
+	end
+
+	-- Build a new context table on top of the existing one, then apply it to
+	-- our virtualNode
+	local existing = virtualNode.context
+	virtualNode.context = assign({}, existing, { [key] = value })
+end
+
 --[[
 	Performs property validation if the static method validateProps is declared.
 	validateProps should follow assert's expected arguments:
@@ -273,7 +313,7 @@ function Component:__mount(reconciler, virtualNode)
 
 	instance.props = props
 
-	local newContext = assign({}, virtualNode.context)
+	local newContext = assign({}, virtualNode.legacyContext)
 	instance._context = newContext
 
 	instance.state = assign({}, instance:__getDerivedState(instance.props, {}))
@@ -284,7 +324,7 @@ function Component:__mount(reconciler, virtualNode)
 	end
 
 	-- It's possible for init() to redefine _context!
-	virtualNode.context = instance._context
+	virtualNode.legacyContext = instance._context
 
 	internalData.lifecyclePhase = ComponentLifecyclePhase.Render
 	local renderResult = instance:render()

src/Component.spec/context.spec.lua

@@ -3,16 +3,17 @@ return function()
 	local createElement = require(script.Parent.Parent.createElement)
 	local createReconciler = require(script.Parent.Parent.createReconciler)
 	local NoopRenderer = require(script.Parent.Parent.NoopRenderer)
+	local oneChild = require(script.Parent.Parent.oneChild)
 
 	local Component = require(script.Parent.Parent.Component)
 
 	local noopReconciler = createReconciler(NoopRenderer)
 
-	it("should be provided as a mutable self._context in Component:init", function()
+	it("should be provided as an internal api on Component", function()
 		local Provider = Component:extend("Provider")
 
 		function Provider:init()
-			self._context.foo = "bar"
+			self:__addContext("foo", "bar")
 		end
 
 		function Provider:render()
@@ -35,7 +36,10 @@ return function()
 
 		local capturedContext
 		function Consumer:init()
-			capturedContext = self._context
+			capturedContext = {
+				hello = self:__getContext("hello"),
+				value = self:__getContext("value"),
+			}
 		end
 
 		function Consumer:render()
@@ -67,7 +71,10 @@ return function()
 
 		local capturedContext
 		function Consumer:init()
-			capturedContext = self._context
+			capturedContext = {
+				hello = self:__getContext("hello"),
+				value = self:__getContext("value"),
+			}
 		end
 
 		function Consumer:render()
@@ -92,12 +99,85 @@ return function()
 		assertDeepEqual(capturedContext, context)
 	end)
 
+	it("should not copy the context table if it doesn't need to", function()
+		local Parent = Component:extend("Parent")
+
+		function Parent:init()
+			self:__addContext("parent", "I'm here!")
+		end
+
+		function Parent:render()
+			-- Create some child element
+			return createElement(function() end)
+		end
+
+		local element = createElement(Parent)
+		local hostParent = nil
+		local hostKey = "Parent"
+		local parentNode = noopReconciler.mountVirtualNode(element, hostParent, hostKey)
+
+		local expectedContext = {
+			parent = "I'm here!",
+		}
+
+		assertDeepEqual(parentNode.context, expectedContext)
+
+		local childNode = oneChild(parentNode.children)
+
+		-- Parent and child should have the same context table
+		expect(parentNode.context, childNode.context)
+	end)
+
+	it("should not allow context to move up the tree", function()
+		local ChildProvider = Component:extend("ChildProvider")
+
+		function ChildProvider:init()
+			self:__addContext("child", "I'm here too!")
+		end
+
+		function ChildProvider:render()
+		end
+
+		local ParentProvider = Component:extend("ParentProvider")
+
+		function ParentProvider:init()
+			self:__addContext("parent", "I'm here!")
+		end
+
+		function ParentProvider:render()
+			return createElement(ChildProvider)
+		end
+
+		local element = createElement(ParentProvider)
+		local hostParent = nil
+		local hostKey = "Parent"
+
+		local parentNode = noopReconciler.mountVirtualNode(element, hostParent, hostKey)
+		local childNode = oneChild(parentNode.children)
+
+		local expectedParentContext = {
+			parent = "I'm here!",
+			-- Context does not travel back up
+		}
+
+		local expectedChildContext = {
+			parent = "I'm here!",
+			child = "I'm here too!"
+		}
+
+		assertDeepEqual(parentNode.context, expectedParentContext)
+		assertDeepEqual(childNode.context, expectedChildContext)
+	end)
+
 	it("should contain values put into the tree by parent nodes", function()
 		local Consumer = Component:extend("Consumer")
 
 		local capturedContext
 		function Consumer:init()
-			capturedContext = self._context
+			capturedContext = {
+				dont = self:__getContext("dont"),
+				frob = self:__getContext("frob"),
+			}
 		end
 
 		function Consumer:render()
@@ -106,7 +186,7 @@ return function()
 		local Provider = Component:extend("Provider")
 
 		function Provider:init()
-			self._context.frob = "ulator"
+			self:__addContext("frob", "ulator")
 		end
 
 		function Provider:render()
@@ -143,11 +223,19 @@ return function()
 	it("should transfer context to children that are replaced", function()
 		local ConsumerA = Component:extend("ConsumerA")
 
+		local function captureAllContext(component)
+			return {
+				A = component:__getContext("A"),
+				B = component:__getContext("B"),
+				frob = component:__getContext("frob"),
+			}
+		end
+
 		local capturedContextA
 		function ConsumerA:init()
-			self._context.A = "hello"
+			self:__addContext("A", "hello")
 
-			capturedContextA = self._context
+			capturedContextA = captureAllContext(self)
 		end
 
 		function ConsumerA:render()
@@ -157,9 +245,9 @@ return function()
 
 		local capturedContextB
 		function ConsumerB:init()
-			self._context.B = "hello"
+			self:__addContext("B", "hello")
 
-			capturedContextB = self._context
+			capturedContextB = captureAllContext(self)
 		end
 
 		function ConsumerB:render()
@@ -168,7 +256,7 @@ return function()
 		local Provider = Component:extend("Provider")
 
 		function Provider:init()
-			self._context.frob = "ulator"
+			self:__addContext("frob", "ulator")
 		end
 
 		function Provider:render()