MockTree.lua

-- MockTree
-- ContextLost
-- April 25th, 2020

local MockTree = {}
MockTree.__index = MockTree

function MockTree:AddJointEdge(joint, part0, part1)
	local joints = self.Joints
	local edges = joints[part0]
	
	if not edges then
		edges = {}
		joints[part0] = edges
	end
	
	local edge = { joint, part1 }
	table.insert(edges, edge)
end

function MockTree:Expand(part0)
	local joints = self.Joints
	local edges = joints[part0]
	
	if not edges then
		return
	end
	
	-- We only want to iterate over part's edges once, remove edges to mark as visited
	joints[part0] = nil
	
	for _,edge in ipairs(edges) do
		local joint, part1 = unpack(edge)
		
		-- Checks if we've already included this part. This will at least be a list with edge back
		-- to parent unless we've already visited this part through another joint and removed it.
		-- Breaks cycles and prioritizes shortest path to root.
		
		if joints[part1] then
			-- Add the parent-child joint edge to the tree list
			table.insert(self.Edges, edge)
			
			-- Recursively add child's edges, DFS order. BFS would
			-- have been fine too, but either works just as well.
			self:Expand(part1)
		end
	end
	
	return self.Edges
end

--[[
	 Returns a list of assembly edges in some tree-sorted order that can be used by `applyTree` to
	 position other parts in `model` relative to `rootPart` if they would be in the same Assembly
	 under a `WorldRoot`. This roughly imitates what the internal spanning tree that `WorldRoot` uses
	 to build an internal transform hierarchy of parts in an assembly, with some limitations:
	
	 - Only supports Motor6D, and Weld. Didn't bother with legacy Motor, Snap, ManualWeld.
	 - Doesn't support Motor/Motor6D.CurrentAngle and co.
	 - Doesn't support WeldConstraints. Can't. Transform isn't exposed to Lua.
	 - Doesn't prioritize by joint type. Weld should take priority over Motor.
	 - Doesn't prioritize by joint/part GUID. Can't. Not exposed to Lua.
	
	 For a resonable model, like an R15 character, that doesn't have duplicate or unsupported joints
	 it should produce the same results as the Roblox spanning tree when applied.
	
	 { { joint, childPart }, ... }
]]--

function MockTree:BuildTree(model, rootPart)
	if not rootPart then
		rootPart = model.PrimaryPart
	end
	
	local tree = MockTree:Construct
	{
		Edges  = {};
		Joints = {};
	}
	
	-- Gather the part-joint graph.
	for _,desc in ipairs(model:GetDescendants()) do
		if desc:IsA("JointInstance") and desc.Enabled then
			local p0 = desc.Part0
			local p1 = desc.Part1
			
			if p0 and p1 then
				-- Add edge to both parts. Assembly joints are bidirectional.
				tree:AddJointEdge(desc, p0, p1)
				tree:AddJointEdge(desc, p1, p0)
			end
		end
	end
	
	-- Build the tree, in order, by recursively following edges out from the root part
	-- Joint edge list map: { [part] = { { joint, otherPart }, ...}, ... }
	
	return tree:Expand(rootPart)
end

function MockTree:ApplyTree(tree)
	for _,edge in ipairs(tree) do
		local joint, childPart = unpack(edge)
		
		local p0 = joint.Part0
		local p1 = joint.Part1
		
		local c0 = joint.C0
		local c1 = joint.C1
		
		if joint:IsA("Motor6D") then
			-- Motor6D, including Motor6D.Transform. Motor6D is now consistently P0->Transform->P1 after recent change.
			local transform = joint.Transform
			
			if p1 == childPart then
				p1.CFrame = p0.CFrame * c0 * transform * c1:Inverse()
			else
				p0.CFrame = p1.CFrame * c1 * transform:Inverse() * c0:Inverse()
			end
		else
			-- Weld
			if p1 == childPart then
				p1.CFrame = p0.CFrame * c0 * c1:Inverse()
			else
				p0.CFrame = p1.CFrame * c1 * c0:Inverse()
			end
		end
	end
end

return MockTree