Fix PatchTree performance issues (#755)

When building the tree, I've implemented a few improvements: - We no longer traverse the full ancestry for every leaf node- we exit early when we find a node that already exists - We no longer search the entire tree to see if a node id exists before creating one with that id, we just check if is in the map
rojo-rbx · Aug 1, 2023 · 34024d8 · 34024d8
1 parent ecc31de
commit 34024d8
Show file tree

Hide file tree

Showing 2 changed files with 92 additions and 59 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -22,6 +22,7 @@
 * Added better support for `Font` properties ([#731])
 * Add new plugin template to the `init` command ([#738])
 * Added rich Source diffs in patch visualizer ([#748])
+* Fix PatchTree performance issues ([#755])
 
 [#745]: https://github.com/rojo-rbx/rojo/pull/745
 [#668]: https://github.com/rojo-rbx/rojo/pull/668
@@ -44,6 +45,7 @@
 [#731]: https://github.com/rojo-rbx/rojo/pull/731
 [#738]: https://github.com/rojo-rbx/rojo/pull/738
 [#748]: https://github.com/rojo-rbx/rojo/pull/748
+[#755]: https://github.com/rojo-rbx/rojo/pull/755
 
 ## [7.3.0] - April 22, 2023
 * Added `$attributes` to project format. ([#574])

diff --git a/plugin/src/PatchTree.lua b/plugin/src/PatchTree.lua
@@ -64,18 +64,18 @@ local Tree = {}
 Tree.__index = Tree
 
 function Tree.new()
-    local tree = {
-        idToNode = {},
+	local tree = {
+		idToNode = {},
 		ROOT = {
 			className = "DataModel",
 			name = "ROOT",
 			children = {},
 		},
-    }
-    -- Add ROOT to idToNode or it won't be found by getNode since that searches *within* ROOT
+	}
+	-- Add ROOT to idToNode or it won't be found by getNode since that searches *within* ROOT
 	tree.idToNode["ROOT"] = tree.ROOT
 
-    return setmetatable(tree, Tree)
+	return setmetatable(tree, Tree)
 end
 
 -- Iterates over all sub-nodes, depth first
@@ -94,23 +94,27 @@ end
 -- Finds a node by id, depth first
 -- searchNode is the node to start the search within, defaults to root
 function Tree:getNode(id, searchNode)
-    if self.idToNode[id] then
-        return self.idToNode[id]
-    end
+	if self.idToNode[id] then
+		return self.idToNode[id]
+	end
 
 	local searchChildren = (searchNode or self.ROOT).children
-    for nodeId, node in searchChildren do
-        if nodeId == id then
-            self.idToNode[id] = node
-            return node
-        end
-        local descendant = self:getNode(id, node)
-        if descendant then
-            return descendant
-        end
-    end
-
-    return nil
+	for nodeId, node in searchChildren do
+		if nodeId == id then
+			self.idToNode[id] = node
+			return node
+		end
+		local descendant = self:getNode(id, node)
+		if descendant then
+			return descendant
+		end
+	end
+
+	return nil
+end
+
+function Tree:doesNodeExist(id)
+	return self.idToNode[id] ~= nil
 end
 
 -- Adds a node to the tree as a child of the node with id == parent
@@ -120,52 +124,53 @@ end
 function Tree:addNode(parent, props)
 	assert(props.id, "props must contain id")
 
-    parent = parent or "ROOT"
+	parent = parent or "ROOT"
 
-    local node = self:getNode(props.id)
-    if node then
+	if self:doesNodeExist(props.id) then
 		-- Update existing node
-        for k, v in props do
-            node[k] = v
-        end
-        return node
-    end
-
-    node = table.clone(props)
-    node.children = {}
+		local node = self:getNode(props.id)
+		for k, v in props do
+			node[k] = v
+		end
+		return node
+	end
+
+	local node = table.clone(props)
+	node.children = {}
 	node.parentId = parent
 
-    local parentNode = self:getNode(parent)
-    if not parentNode then
-        Log.warn("Failed to create node since parent doesnt exist: {}, {}", parent, props)
-        return
-    end
+	local parentNode = self:getNode(parent)
+	if not parentNode then
+		Log.warn("Failed to create node since parent doesnt exist: {}, {}", parent, props)
+		return
+	end
 
-    parentNode.children[node.id] = node
-    self.idToNode[node.id] = node
+	parentNode.children[node.id] = node
+	self.idToNode[node.id] = node
 
-    return node
+	return node
 end
 
 -- Given a list of ancestor ids in descending order, builds the nodes for them
 -- using the patch and instanceMap info
-function Tree:buildAncestryNodes(ancestryIds: { string }, patch, instanceMap)
-    -- Build nodes for ancestry by going up the tree
-    local previousId = "ROOT"
-    for _, ancestorId in ancestryIds do
-        local value = instanceMap.fromIds[ancestorId] or patch.added[ancestorId]
-        if not value then
-            Log.warn("Failed to find ancestor object for " .. ancestorId)
-            continue
-        end
-        self:addNode(previousId, {
-            id = ancestorId,
-            className = value.ClassName,
-            name = value.Name,
-            instance = if typeof(value) == "Instance" then value else nil,
-        })
-        previousId = ancestorId
-    end
+function Tree:buildAncestryNodes(previousId: string?, ancestryIds: { string }, patch, instanceMap)
+	-- Build nodes for ancestry by going up the tree
+	previousId = previousId or "ROOT"
+
+	for _, ancestorId in ancestryIds do
+		local value = instanceMap.fromIds[ancestorId] or patch.added[ancestorId]
+		if not value then
+			Log.warn("Failed to find ancestor object for " .. ancestorId)
+			continue
+		end
+		self:addNode(previousId, {
+			id = ancestorId,
+			className = value.ClassName,
+			name = value.Name,
+			instance = if typeof(value) == "Instance" then value else nil,
+		})
+		previousId = ancestorId
+	end
 end
 
 local PatchTree = {}
@@ -175,6 +180,8 @@ local PatchTree = {}
 function PatchTree.build(patch, instanceMap, changeListHeaders)
 	local tree = Tree.new()
 
+	local knownAncestors = {}
+
 	for _, change in patch.updated do
 		local instance = instanceMap.fromIds[change.id]
 		if not instance then
@@ -185,13 +192,21 @@ function PatchTree.build(patch, instanceMap, changeListHeaders)
 		local ancestryIds = {}
 		local parentObject = instance.Parent
 		local parentId = instanceMap.fromInstances[parentObject]
+		local previousId = nil
 		while parentObject do
+			if knownAncestors[parentId] then
+				-- We've already added this ancestor
+				previousId = parentId
+				break
+			end
+
 			table.insert(ancestryIds, 1, parentId)
+			knownAncestors[parentId] = true
 			parentObject = parentObject.Parent
 			parentId = instanceMap.fromInstances[parentObject]
 		end
 
-		tree:buildAncestryNodes(ancestryIds, patch, instanceMap)
+		tree:buildAncestryNodes(previousId, ancestryIds, patch, instanceMap)
 
 		-- Gather detail text
 		local changeList, hint = nil, nil
@@ -268,14 +283,22 @@ function PatchTree.build(patch, instanceMap, changeListHeaders)
 		local ancestryIds = {}
 		local parentObject = instance.Parent
 		local parentId = instanceMap.fromInstances[parentObject] or HttpService:GenerateGUID(false)
+		local previousId = nil
 		while parentObject do
+			if knownAncestors[parentId] then
+				-- We've already added this ancestor
+				previousId = parentId
+				break
+			end
+
 			instanceMap:insert(parentId, parentObject) -- This ensures we can find the parent later
 			table.insert(ancestryIds, 1, parentId)
+			knownAncestors[parentId] = true
 			parentObject = parentObject.Parent
 			parentId = instanceMap.fromInstances[parentObject] or HttpService:GenerateGUID(false)
 		end
 
-		tree:buildAncestryNodes(ancestryIds, patch, instanceMap)
+		tree:buildAncestryNodes(previousId, ancestryIds, patch, instanceMap)
 
 		-- Add this node to tree
 		local nodeId = instanceMap.fromInstances[instance] or HttpService:GenerateGUID(false)
@@ -295,8 +318,16 @@ function PatchTree.build(patch, instanceMap, changeListHeaders)
 		local parentId = change.Parent
 		local parentData = patch.added[parentId]
 		local parentObject = instanceMap.fromIds[parentId]
+		local previousId = nil
 		while parentId do
+			if knownAncestors[parentId] then
+				-- We've already added this ancestor
+				previousId = parentId
+				break
+			end
+
 			table.insert(ancestryIds, 1, parentId)
+			knownAncestors[parentId] = true
 			parentId = nil
 
 			if parentData then
@@ -312,7 +343,7 @@ function PatchTree.build(patch, instanceMap, changeListHeaders)
 			end
 		end
 
-		tree:buildAncestryNodes(ancestryIds, patch, instanceMap)
+		tree:buildAncestryNodes(previousId, ancestryIds, patch, instanceMap)
 
 		-- Gather detail text
 		local changeList, hint = nil, nil