From 8b17d3b7d9c251af3f281b92f0d8b7e7bafda6b2 Mon Sep 17 00:00:00 2001
From: Lucien Greathouse <me@lpghatguy.com>
Date: Mon, 21 May 2018 12:48:25 -0700
Subject: [PATCH] Intense robustness pass

---
 plugin/src/Plugin.lua     |  4 ++-
 plugin/src/Reconciler.lua | 52 +++++++++++++++++++--------------------
 plugin/src/RouteMap.lua   | 38 ++++++++++++++++++++--------
 3 files changed, 57 insertions(+), 37 deletions(-)

diff --git a/plugin/src/Plugin.lua b/plugin/src/Plugin.lua
index fa0f4c72..498f6be3 100644
--- a/plugin/src/Plugin.lua
+++ b/plugin/src/Plugin.lua
@@ -66,7 +66,9 @@ end
 function Plugin:restart()
 	warn("Rojo: The server has changed since the last request, reloading plugin...")
 
-	self._reconciler:clear()
+	self._reconciler:destruct()
+	self._reconciler = Reconciler.new()
+
 	self._api = nil
 	self._polling = false
 	self._syncInProgress = false
diff --git a/plugin/src/Reconciler.lua b/plugin/src/Reconciler.lua
index a0ea06a8..327aa839 100644
--- a/plugin/src/Reconciler.lua
+++ b/plugin/src/Reconciler.lua
@@ -1,13 +1,26 @@
 local RouteMap = require(script.Parent.RouteMap)
 
-local function classEqual(rbx, className)
-	if className == "*" then
+local function classEqual(a, b)
+	if a == "*" or b == "*" then
 		return true
 	end
 
-	return rbx.ClassName == className
+	return a == b
 end
 
+local function applyProperties(target, properties)
+	for key, property in pairs(properties) do
+		-- TODO: Transform property value based on property.Type
+		-- Right now, we assume that 'value' is primitive!
+		target[key] = property.Value
+	end
+end
+
+--[[
+	Attempt to parent `rbx` to `parent`, doing nothing if:
+	* parent is already `parent`
+	* Changing parent threw an error
+]]
 local function reparent(rbx, parent)
 	if rbx then
 		if rbx.Parent == parent then
@@ -42,7 +55,7 @@ local function findNextChildPair(primaryChildren, secondaryChildren, visited)
 			visited[primaryChild] = true
 
 			for _, secondaryChild in ipairs(secondaryChildren) do
-				if primaryChild.ClassName == secondaryChild.ClassName and primaryChild.Name == secondaryChild.Name then
+				if classEqual(primaryChild.ClassName, secondaryChild.ClassName) and primaryChild.Name == secondaryChild.Name then
 					visited[secondaryChild] = true
 
 					return primaryChild, secondaryChild
@@ -114,12 +127,11 @@ function Reconciler:_reify(item)
 	local rbx = Instance.new(className)
 	rbx.Name = item.Name
 
-	for key, property in pairs(item.Properties) do
-		-- TODO: Check for compound types, like Vector3!
-		rbx[key] = property.Value
-	end
+	applyProperties(rbx, item.Properties)
 
-	self:_reconcileChildren(rbx, item)
+	for _, child in ipairs(item.Children) do
+		reparent(self:_reify(child), rbx)
+	end
 
 	if item.Route then
 		self._routeMap:insert(item.Route, rbx)
@@ -129,10 +141,10 @@ function Reconciler:_reify(item)
 end
 
 --[[
-	Clears any state that the Reconciler has, effectively restarting it.
+	Clears any state that the Reconciler has, stopping it completely.
 ]]
-function Reconciler:clear()
-	self._routeMap:clear()
+function Reconciler:destruct()
+	self._routeMap:destruct()
 end
 
 --[[
@@ -156,28 +168,16 @@ function Reconciler:reconcile(rbx, item)
 	end
 
 	-- Item changed type!
-	if not classEqual(rbx, item.ClassName) then
+	if not classEqual(rbx.ClassName, item.ClassName) then
 		self._routeMap:removeByRbx(rbx)
 		rbx:Destroy()
 
 		return self:_reify(item)
 	end
 
-	-- Apply all properties, Roblox will de-duplicate changes
-	for key, property in pairs(item.Properties) do
-		-- TODO: Transform property value based on property.Type
-		-- Right now, we assume that 'value' is primitive!
-
-		rbx[key] = property.Value
-	end
-
-	-- Use a dumb algorithm for reconciling children
+	applyProperties(rbx, item.Properties)
 	self:_reconcileChildren(rbx, item)
 
-	if item.Route then
-		self._routeMap:insert(item.Route, rbx)
-	end
-
 	return rbx
 end
 
diff --git a/plugin/src/RouteMap.lua b/plugin/src/RouteMap.lua
index 7b385ce1..35d67ff6 100644
--- a/plugin/src/RouteMap.lua
+++ b/plugin/src/RouteMap.lua
@@ -25,6 +25,10 @@ end
 function RouteMap:insert(route, rbx)
 	local hashed = hashRoute(route)
 
+	-- Make sure that each route and instance are only present in RouteMap once.
+	self:removeByRoute(route)
+	self:removeByRbx(rbx)
+
 	self._map[hashed] = rbx
 	self._reverseMap[rbx] = hashed
 	self._connectionsByRbx[rbx] = rbx.AncestryChanged:Connect(function(_, parent)
@@ -42,26 +46,36 @@ function RouteMap:removeByRoute(route)
 	local hashedRoute = hashRoute(route)
 	local rbx = self._map[hashedRoute]
 
-	if rbx then
-		self._map[hashedRoute] = nil
-		self._reverseMap[rbx] = nil
-		self._connectionsByRbx[rbx] = nil
+	if rbx ~= nil then
+		self:_removeInternal(hashedRoute, rbx)
 	end
 end
 
 function RouteMap:removeByRbx(rbx)
 	local hashedRoute = self._reverseMap[rbx]
 
-	if hashedRoute then
-		self._map[hashedRoute] = nil
-		self._reverseMap[rbx] = nil
-		self._connectionsByRbx[rbx] = nil
+	if hashedRoute ~= nil then
+		self:_removeInternal(hashedRoute, rbx)
 	end
+end
+
+--[[
+	Correcly removes the given Roblox Instance/Route pair from the RouteMap.
+]]
+function RouteMap:_removeInternal(rbx, hashedRoute)
+	self._map[hashedRoute] = nil
+	self._reverseMap[rbx] = nil
+	self._connectionsByRbx[rbx]:Disconnect()
+	self._connectionsByRbx[rbx] = nil
 
 	self:removeRbxDescendants(rbx)
 end
 
-function RouteMap:removeRbxDescendants(parentRbx)
+--[[
+	Ensure that there are no descendants of the given Roblox Instance still
+	present in the map, guaranteeing that it has been cleaned out.
+]]
+function RouteMap:_removeRbxDescendants(parentRbx)
 	for rbx in pairs(self._reverseMap) do
 		if rbx:IsDescendantOf(parentRbx) then
 			self:removeByRbx(rbx)
@@ -69,7 +83,11 @@ function RouteMap:removeRbxDescendants(parentRbx)
 	end
 end
 
-function RouteMap:clear()
+--[[
+	Remove all items from the map and disconnect all connections, cleaning up
+	the RouteMap.
+]]
+function RouteMap:destruct()
 	self._map = {}
 	self._reverseMap = {}