MockTree: A module for handling Welds and Motor6Ds outside of Workspace

MockTree.lua

local MockTree = {}

local function addJointEdge(joints, joint, me, other)
	local edgeList = joints[me]
	if not edgeList then
		edgeList = {}
		joints[me] = edgeList
	end
	table.insert(edgeList, {joint, other})
end

local function expandTree(tree, joints, parentPart, parentJoints)
	if not parentJoints then
		return
	end
	-- We only want to iterate over part's edges once, remove edges to mark as visited
	joints[parentPart] = nil
	
	for _, edge in ipairs(parentJoints) do
		local joint, childPart = unpack(edge)
		local childJoints = joints[childPart]
		-- 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.
		-- Breacks cycles and prioritizes shortest path to root.
		if childJoints then
			-- Add the parent-child joint edge to the tree list
			table.insert(tree, edge)
			-- Recursively add child's edges, DFS order. BFS would have been fine too, but either
			-- works just as well.
			expandTree(tree, joints, childPart, childJoints)
		end
	end
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)
	-- Joint edge list map: { [part] = { { joint, otherPart }, ...}, ... }
	local joints = {}
	-- Gather the part-joint graph.
	for _, child in ipairs(model:GetDescendants()) do
		if child:IsA("Motor6D") or child:IsA("Weld") and child.Enabled then
			local p0 = child.Part0
			local p1 = child.Part1
			if p0 and p1 and child.Enabled then
				-- Add edge to both parts. Assembly joints are bidirectional.
				addJointEdge(joints, child, p0, p1)
				addJointEdge(joints, child, p1, p0)
			end
		end
	end
	
	local tree = {}
	-- Build the tree, in order, by recursively following edges out from the root part
	expandTree(tree, joints, rootPart, joints[rootPart])
	
	return tree
end

-- Applies a tree generated by `buildTree`, positioning other parts relative to the trees root part.
function MockTree.applyTree(tree)
	for _, edge in ipairs(tree) do
		local joint, childPart = unpack(edge)
		local p0 = joint.Part0
		local p1 = joint.Part1
		if joint:IsA("Motor6D") then
			-- Motor6D, including Motor6D.Transform. Motor6D is now consistently P0->Transform->P1 after recent change.
			if p1 == childPart then
				p1.CFrame = p0.CFrame * joint.C0 * joint.Transform * joint.C1:Inverse()
			else
				p0.CFrame = p1.CFrame * joint.C1 * joint.Transform:Inverse() * joint.C0:Inverse()
			end
		else
			-- Weld
			if p1 == childPart then
				p1.CFrame = p0.CFrame * joint.C0 * joint.C1:Inverse()
			else
				p0.CFrame = p1.CFrame * joint.C1 * joint.C0:Inverse()
			end
		end
	end
end

return MockTree