Flex.lua

--[[
	A special key for property tables.
	
	When applied to a UIListLayout, enables flexible sizing on the laid out UI
	objects.
	
	When applied to a GuiObject in a [Flex]-ed UIListLayout, opts the object
	into flexible sizing - the object's size will be overwritten along the
	direction of layout to fill up spare space.
	
	TODO: make this work with UISizeConstraint
	TODO: make this work with padding
	TODO: optimise away some unnecessary recalculateFlexSize calls
]]

local Flex = {}
Flex.type = "SpecialKey"
Flex.kind = "Flex"
Flex.stage = "observer"

local flexObjectSet = setmetatable({}, {__mode = "k"})

local function recalculateFlexSizes(layout)
	local parent = layout.Parent
	if parent == nil or not parent:IsA("GuiBase2d") then
		return
	end
	local children = parent:GetChildren()
	
	local layoutAxisIsX = layout.FillDirection == Enum.FillDirection.Horizontal
	local layoutAxis = if layoutAxisIsX then Vector2.xAxis else Vector2.yAxis
	
	local flexes = {}
	local numFlexes = 0
	local collapsedLength = 0
	for _, child in children do
		if flexObjectSet[child] ~= nil then
			table.insert(flexes, child)
			numFlexes += 1
		elseif child:IsA("GuiObject") then
			collapsedLength += if layoutAxisIsX then child.AbsoluteSize.X else child.AbsoluteSize.Y
		end
	end

	if numFlexes == 0 then
		return
	end

	local desiredSizeProperty = if parent:IsA("ScrollingFrame") then "AbsoluteCanvasSize" else "AbsoluteSize"
	local desiredLength = if layoutAxisIsX then parent[desiredSizeProperty].X else parent[desiredSizeProperty].Y
	local spareSpace = math.max(0, desiredLength - collapsedLength)
	local filledSpace = 0
	for index, flex in flexes do
		local idealLength = (spareSpace - filledSpace) / (numFlexes - index + 1)
		local realLength = math.round(idealLength)
		filledSpace += realLength
		flex.Size = UDim2.new(
			if layoutAxisIsX then UDim.new(0, realLength) else flex.Size.X,
			if not layoutAxisIsX then UDim.new(0, realLength) else flex.Size.Y
		)
	end
end

function Flex:apply(_, applyToRef, cleanupTasks)
	local instance = applyToRef.instance :: Instance

	if instance:IsA("UIListLayout") then
		recalculateFlexSizes(instance)
		local parentConnections = {}
		table.insert(cleanupTasks, {
			parentConnections,
			instance.AncestryChanged:Connect(function()
				recalculateFlexSizes(instance)

				for _, conn in parentConnections do
					conn:Disconnect()
				end
				table.clear(parentConnections)

				local parent = instance.Parent

				if parent ~= nil and parent:IsA("GuiBase2d") then
					local desiredContentSizeProperty = if parent:IsA("ScrollingFrame") then "AbsoluteCanvasSize" else "AbsoluteSize"

					table.insert(parentConnections, 
						parent:GetPropertyChangedSignal(desiredContentSizeProperty):Connect(function()
							recalculateFlexSizes(instance)
						end)
					)
					table.insert(parentConnections, 
						parent.ChildAdded:Connect(function(child)
							if flexObjectSet[child] ~= nil then
								recalculateFlexSizes(instance)
							end
						end)
					)
					table.insert(parentConnections, 
						parent.ChildRemoved:Connect(function(child)
							if flexObjectSet[child] ~= nil then
								recalculateFlexSizes(instance)
							end
						end)
					)
					
					if parent:IsA("GuiObject") then
						local controlSizeProperty = if parent:IsA("ScrollingFrame") then "CanvasSize" else "Size"
						
						table.insert(parentConnections, 
							parent:GetPropertyChangedSignal(controlSizeProperty):Connect(function()
								recalculateFlexSizes(instance)
							end)
						)
					end
				end

			end)
		})
	elseif instance:IsA("GuiObject") then
		flexObjectSet[instance] = true
		table.insert(cleanupTasks, function()
			flexObjectSet[instance] = nil
		end)
	end
end

return Flex