rojo/plugin/src/PatchTree.lua

--[[
	Methods to turn PatchSets into trees matching the DataModel containing
	the changes and metadata for use in the PatchVisualizer component.
]]

local HttpService = game:GetService("HttpService")
local ScriptEditorService = game:GetService("ScriptEditorService")

local Rojo = script:FindFirstAncestor("Rojo")
local Plugin = Rojo.Plugin
local Packages = Rojo.Packages

local Log = require(Packages.Log)

local ChangeMetadata = require(Plugin.ChangeMetadata)
local Config = require(Plugin.Config)
local SHA1 = require(Plugin.SHA1)
local Timer = require(Plugin.Timer)
local Types = require(Plugin.Types)
local decodeValue = require(Plugin.Reconciler.decodeValue)
local getProperty = require(Plugin.Reconciler.getProperty)

local function yieldIfNeeded(clock)
	if os.clock() - clock > 1 / 20 then
		task.wait()
		return os.clock()
	end
	return clock
end

local function alphabeticalNext(t, state)
	-- Equivalent of the next function, but returns the keys in the alphabetic
	-- order of node names. We use a temporary ordered key table that is stored in the
	-- table being iterated.

	local key = nil
	if state == nil then
		-- First iteration, generate the index
		local orderedIndex, i = table.create(5), 0
		for k in t do
			i += 1
			orderedIndex[i] = k
		end
		table.sort(orderedIndex, function(a, b)
			local nodeA, nodeB = t[a], t[b]
			return (nodeA.name or "") < (nodeB.name or "")
		end)

		t.__orderedIndex = orderedIndex
		key = orderedIndex[1]
	else
		-- Fetch the next value
		for i, orderedState in t.__orderedIndex do
			if orderedState == state then
				key = t.__orderedIndex[i + 1]
				break
			end
		end
	end

	if key then
		return key, t[key]
	end

	-- No more value to return, cleanup
	t.__orderedIndex = nil

	return
end

local function alphabeticalPairs(t)
	-- Equivalent of the pairs() iterator, but sorted
	return alphabeticalNext, t, nil
end

local Tree = {}
Tree.__index = Tree

function Tree.new()
	local tree = {
		idToNode = {},
		ROOT = {
			id = "ROOT",
			className = "DataModel",
			name = "game",
			children = {},
		},
	}
	-- 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)
end

-- Iterates over all nodes and counts them up
function Tree:getCount()
	local count = 0
	self:forEach(function()
		count += 1
	end)
	return count
end

-- Iterates over all sub-nodes, depth first
-- node is where to start from, defaults to root
-- depth is used for recursion but can be used to set the starting depth
function Tree:forEach(callback, node, depth)
	depth = depth or 1
	for _, child in alphabeticalPairs(if node then node.children else self.ROOT.children) do
		callback(child, depth)
		if type(child.children) == "table" then
			self:forEach(callback, child, depth + 1)
		end
	end
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

	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
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
-- If parent is nil, it defaults to root
-- props must contain id, and cannot contain children or parentId
-- other than those three, it can hold anything
function Tree:addNode(parent, props)
	assert(props.id, "props must contain id")

	parent = parent or "ROOT"

	if self:doesNodeExist(props.id) then
		-- Update existing node
		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

	parentNode.children[node.id] = node
	self.idToNode[node.id] = 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(previousId: string?, ancestryIds: { string }, patch, instanceMap)
	local clock = os.clock()
	-- Build nodes for ancestry by going up the tree
	previousId = previousId or "ROOT"

	for _, ancestorId in ancestryIds do
		clock = yieldIfNeeded(clock)

		local value = instanceMap.fromIds[ancestorId] or patch.added[ancestorId]
		if not value then
			-- Ancestor is missing from both instanceMap and patch.added.
			-- This can happen during rapid changes (e.g., undo operations) where
			-- patches arrive referencing instances that were never fully synced.
			-- Create a placeholder node to maintain the ancestry chain, allowing
			-- the visualization to still work with incomplete information.
			Log.warn("Failed to find ancestor object for {}; creating placeholder node", ancestorId)
			self:addNode(previousId, {
				id = ancestorId,
				className = "Unknown",
				name = "[Unknown Ancestor]",
				instance = nil,
			})
			previousId = 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 function computeRootChangeInfo(tree)
	local totalLinesAdded = 0
	local totalLinesRemoved = 0
	local totalPropChanges = 0
	local hasLineData = false

	tree:forEach(function(node)
		if not node.changeInfo then
			return
		end
		if node.changeInfo.linesAdded then
			totalLinesAdded += node.changeInfo.linesAdded
			hasLineData = true
		end
		if node.changeInfo.linesRemoved then
			totalLinesRemoved += node.changeInfo.linesRemoved
			hasLineData = true
		end
		if node.changeInfo.propChanges then
			totalPropChanges += node.changeInfo.propChanges
		end
	end)

	if not hasLineData and totalPropChanges == 0 then
		return nil
	end

	return {
		linesAdded = if hasLineData then totalLinesAdded else nil,
		linesRemoved = if hasLineData then totalLinesRemoved else nil,
		isWhitespaceOnly = false,
		propChanges = if totalPropChanges > 0 then totalPropChanges else nil,
	}
end

local PatchTree = {}

-- Builds a new tree from a patch and instanceMap
-- uses changeListHeaders in node.changeList
-- gitMetadata is optional: { changedIds: {string}, scriptCommittedHashes: {[string]: {string}} }
function PatchTree.build(patch, instanceMap, changeListHeaders, gitMetadata)
	Timer.start("PatchTree.build")
	local clock = os.clock()

	local tree = Tree.new()

	local knownAncestors = {}

	local changedIdSet = nil
	if gitMetadata and gitMetadata.changedIds then
		changedIdSet = {}
		for _, id in gitMetadata.changedIds do
			changedIdSet[id] = true
		end
	end

	Timer.start("patch.updated")
	for _, change in patch.updated do
		clock = yieldIfNeeded(clock)

		local instance = instanceMap.fromIds[change.id]
		if not instance then
			continue
		end

		-- Gather ancestors from existing DOM
		local ancestryIds = {}
		local parentObject = instance.Parent
		local parentId = instanceMap.fromInstances[parentObject]
		local previousId = nil
		while parentObject and parentId 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(previousId, ancestryIds, patch, instanceMap)

		-- Gather detail text
		local changeList, changeInfo = nil, nil
		if next(change.changedProperties) or change.changedName then
			changeList = {}

			local changeIndex = 0
			local function addProp(prop: string, current: any?, incoming: any?)
				changeIndex += 1
				changeList[changeIndex] = { prop, current, incoming }
			end

			-- Gather the changes

			if change.changedName then
				addProp("Name", instance.Name, change.changedName)
			end

			for prop, incoming in change.changedProperties do
				local incomingSuccess, incomingValue = decodeValue(incoming, instanceMap)
				local currentSuccess, currentValue = getProperty(instance, prop)

				-- If incoming is a raw number and current is an EnumItem, convert incoming
				-- to the proper EnumItem for better display (server sends Enum as raw numbers)
				if
					incomingSuccess
					and currentSuccess
					and typeof(incomingValue) == "number"
					and typeof(currentValue) == "EnumItem"
				then
					local enumType = currentValue.EnumType
					local ok, converted = pcall(function()
						return enumType:FromValue(incomingValue)
					end)
					if ok and converted then
						incomingValue = converted
					end
				end

				addProp(
					prop,
					if currentSuccess then currentValue else "[Error]",
					if incomingSuccess then incomingValue else select(2, next(incoming))
				)
			end

			-- Compute structured change metadata via ChangeMetadata module
			changeInfo = ChangeMetadata.compute(changeList)

			-- Sort changes and add header
			table.sort(changeList, function(a, b)
				return a[1] < b[1]
			end)
			table.insert(changeList, 1, changeListHeaders)
		end

		-- Compute default selection from git metadata
		local defaultSelection = nil
		if changedIdSet then
			if not changedIdSet[change.id] then
				defaultSelection = "pull"
			elseif instance:IsA("LuaSourceContainer") and gitMetadata.scriptCommittedHashes then
				local hashes = gitMetadata.scriptCommittedHashes[change.id]
				if hashes then
					local source = instance.Source
					local gitBlob = "blob " .. tostring(#source) .. "\0" .. source
					local studioHash = SHA1(buffer.fromstring(gitBlob))
					for _, hash in hashes do
						if studioHash == hash then
							defaultSelection = "push"
							break
						end
					end

					if defaultSelection == nil then
						local ok, draft = pcall(ScriptEditorService.GetEditorSource, ScriptEditorService, instance)
						if ok and draft ~= source then
							local draftBlob = "blob " .. tostring(#draft) .. "\0" .. draft
							local draftHash = SHA1(buffer.fromstring(draftBlob))
							for _, hash in hashes do
								if draftHash == hash then
									defaultSelection = "push"
									break
								end
							end
						end
					end
				end
			end
		end

		tree:addNode(instanceMap.fromInstances[instance.Parent], {
			id = change.id,
			patchType = "Edit",
			className = instance.ClassName,
			name = instance.Name,
			instance = instance,
			changeInfo = changeInfo,
			changeList = changeList,
			defaultSelection = defaultSelection,
		})
	end
	Timer.stop()

	Timer.start("patch.removed")
	for _, idOrInstance in patch.removed do
		clock = yieldIfNeeded(clock)

		local instance = if Types.RbxId(idOrInstance) then instanceMap.fromIds[idOrInstance] else idOrInstance
		if not instance then
			-- If we're viewing a past patch, the instance is already removed
			-- and we therefore cannot get the tree for it anymore
			continue
		end

		-- Skip auto-created instances that shouldn't be synced
		if Config.ignoredClassNames[instance.ClassName] then
			continue
		end

		-- Gather ancestors from existing DOM
		-- (note that they may have no ID if they're being removed as unknown)
		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(previousId, ancestryIds, patch, instanceMap)

		-- Compute line count metadata for removed scripts
		local changeInfo = ChangeMetadata.computeForRemoval(instance)

		-- Add this node to tree
		-- Default to nil (unselected) - user must explicitly choose Push/Pull/Skip
		local nodeId = instanceMap.fromInstances[instance] or HttpService:GenerateGUID(false)
		instanceMap:insert(nodeId, instance)
		tree:addNode(instanceMap.fromInstances[instance.Parent], {
			id = nodeId,
			patchType = "Remove",
			className = instance.ClassName,
			name = instance.Name,
			instance = instance,
			changeInfo = changeInfo,
			defaultSelection = nil,
		})
	end
	Timer.stop()

	Timer.start("patch.added")
	for id, change in patch.added do
		clock = yieldIfNeeded(clock)

		-- Gather ancestors from existing DOM or future additions
		local ancestryIds = {}
		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
				-- object is parented to an instance that does not exist yet
				parentId = parentData.Parent
				parentData = patch.added[parentId]
				parentObject = instanceMap.fromIds[parentId]
			elseif parentObject then
				-- object is parented to an instance that exists
				parentObject = parentObject.Parent
				parentId = instanceMap.fromInstances[parentObject]
				parentData = patch.added[parentId]
			end
		end

		tree:buildAncestryNodes(previousId, ancestryIds, patch, instanceMap)

		-- Gather detail text
		local changeList, changeInfo = nil, nil
		if next(change.Properties) then
			changeList = {}

			local changeIndex = 0
			local function addProp(prop: string, incoming: any)
				changeIndex += 1
				changeList[changeIndex] = { prop, "N/A", incoming }
			end

			for prop, incoming in change.Properties do
				local success, incomingValue = decodeValue(incoming, instanceMap)
				addProp(prop, if success then incomingValue else select(2, next(incoming)))
			end

			-- Compute structured change metadata via ChangeMetadata module
			changeInfo = ChangeMetadata.computeForAddition(changeList)

			-- Sort changes and add header
			table.sort(changeList, function(a, b)
				return a[1] < b[1]
			end)
			table.insert(changeList, 1, changeListHeaders)
		end

		-- Add this node to tree
		-- Default to nil (unselected) - user must explicitly choose Push/Pull/Skip
		tree:addNode(change.Parent, {
			id = change.Id,
			patchType = "Add",
			className = change.ClassName,
			name = change.Name,
			changeInfo = changeInfo,
			changeList = changeList,
			instance = instanceMap.fromIds[id],
			defaultSelection = nil,
		})
	end
	Timer.stop()

	tree.ROOT.changeInfo = computeRootChangeInfo(tree)

	Timer.stop()
	return tree
end

-- Builds initial selections map from the tree using defaultSelection values
-- Returns a table mapping node id -> selection ("push", "pull", or "ignore")
function PatchTree.buildInitialSelections(tree)
	local selections = {}
	tree:forEach(function(node)
		if node.patchType and node.defaultSelection then
			selections[node.id] = node.defaultSelection
		end
	end)
	return selections
end

-- Gets the count of unselected items (nodes with patchType but no selection)
-- Returns count of unselected items
function PatchTree.countUnselected(tree, selections)
	local count = 0
	tree:forEach(function(node)
		if node.patchType and selections[node.id] == nil then
			count += 1
		end
	end)
	return count
end

-- Gets all selectable node IDs (nodes that have a patchType)
-- Returns an array of node IDs
function PatchTree.getSelectableNodeIds(tree)
	local ids = {}
	tree:forEach(function(node)
		if node.patchType then
			table.insert(ids, node.id)
		end
	end)
	return ids
end

-- Checks if all selectable nodes have a selection
function PatchTree.allNodesSelected(tree, selections)
	local allSelected = true
	tree:forEach(function(node)
		if node.patchType and not selections[node.id] then
			allSelected = false
		end
	end)
	return allSelected
end

-- Updates the metadata of a tree with the unapplied patch and currently existing instances
-- Builds a new tree from the data if one isn't provided
-- Always returns a new tree for immutability purposes in Roact
function PatchTree.updateMetadata(tree, patch, instanceMap, unappliedPatch)
	Timer.start("PatchTree.updateMetadata")
	if tree then
		-- A shallow copy is enough for our purposes here since we really only need a new top-level object
		-- for immutable comparison checks in Roact
		tree = table.clone(tree)
	else
		tree = PatchTree.build(patch, instanceMap)
	end

	-- Update isWarning metadata
	Timer.start("isWarning")
	for _, failedChange in unappliedPatch.updated do
		local node = tree:getNode(failedChange.id)
		if not node then
			continue
		end

		node.isWarning = true
		Log.trace("Marked node as warning: {} {}", node.id, node.name)

		if not node.changeList then
			continue
		end

		local warnings = 0
		for _, change in node.changeList do
			local property = change[1]
			local propertyFailedToApply = if property == "Name"
				then failedChange.changedName ~= nil -- Name is not in changedProperties, so it needs a special case
				else failedChange.changedProperties[property] ~= nil

			if not propertyFailedToApply then
				-- This change didn't fail, no need to mark
				continue
			end

			warnings += 1
			if change[4] == nil then
				change[4] = { isWarning = true }
			else
				change[4].isWarning = true
			end
			Log.trace("  Marked property as warning: {}.{}", node.name, property)
		end

		-- Preserve existing metadata fields (linesAdded, etc.) and add failure tracking
		node.changeInfo = table.clone(node.changeInfo)
		node.changeInfo.edits = (node.changeInfo.edits or (#node.changeList - 1)) - warnings
		node.changeInfo.failed = if warnings > 0 then warnings else nil
	end
	for failedAdditionId in unappliedPatch.added do
		local node = tree:getNode(failedAdditionId)
		if not node then
			continue
		end

		node.isWarning = true
		Log.trace("Marked node as warning: {} {}", node.id, node.name)

		if not node.changeList then
			continue
		end

		for _, change in node.changeList do
			-- Failed addition means that all properties failed to be added
			if change[4] == nil then
				change[4] = { isWarning = true }
			else
				change[4].isWarning = true
			end
			Log.trace("  Marked property as warning: {}.{}", node.name, change[1])
		end

		-- Preserve existing metadata fields and mark all edits as failed
		local totalEdits = node.changeInfo.edits or (#node.changeList - 1)
		node.changeInfo = table.clone(node.changeInfo)
		node.changeInfo.failed = totalEdits
		node.changeInfo.edits = nil
	end
	for _, failedRemovalIdOrInstance in unappliedPatch.removed do
		local failedRemovalId = if Types.RbxId(failedRemovalIdOrInstance)
			then failedRemovalIdOrInstance
			else instanceMap.fromInstances[failedRemovalIdOrInstance]
		if not failedRemovalId then
			continue
		end

		local node = tree:getNode(failedRemovalId)
		if not node then
			continue
		end

		node.isWarning = true
		Log.trace("Marked node as warning: {} {}", node.id, node.name)
	end
	Timer.stop()

	-- Update if instances exist
	Timer.start("instanceAncestry")
	tree:forEach(function(node)
		if node.instance then
			if node.instance.Parent == nil and node.instance ~= game then
				-- This instance has been removed
				Log.trace("Removed instance from node: {} {}", node.id, node.name)
				node.instance = nil
			end
		else
			-- This instance may have been added
			node.instance = instanceMap.fromIds[node.id]
			if node.instance then
				Log.trace("Added instance to node: {} {}", node.id, node.name)
			end
		end
	end)
	Timer.stop()

	tree.ROOT.changeInfo = computeRootChangeInfo(tree)

	Timer.stop()
	return tree
end

return PatchTree