A basic R tree

RTree.luau

--!strict
--!native
-- Services
-- Packages
-- Modules
-- Types
type RectangleImpl = {
	__index: RectangleImpl,
	__tostring: (self: Rectangle) -> string,
	__eq: (self: Rectangle, other: Rectangle) -> boolean,
	__lt: (self: Rectangle, other: Rectangle) -> boolean,
	__le: (self: Rectangle, other: Rectangle) -> boolean,
	__len: (self: Rectangle) -> number,
	ToRect: (self: Rectangle, other: Rectangle) -> Rect,
	Contains: (self: Rectangle, other: Rectangle) -> boolean,
	Overlaps: (self: Rectangle, other: Rectangle) -> boolean,
	GetOverlap: (self: Rectangle, other: Rectangle) -> number,
}
export type Rectangle = typeof(setmetatable(
	{} :: {
		_Key: string,
		MinX: number,
		MinY: number,
		MaxX: number,
		MaxY: number,
		Width: number,
		Height: number,
		Area: number,
	},
	{} :: RectangleImpl
))
type Directory<V> = {
	Rectangle: Rectangle,
	Item: V?,
	SubDirectories: { [string]: Directory<V> },
}
type RTreeImpl<V> = {
	__index: RTreeImpl<V>,
	_Add: (self: RTree<V>, rect: Rectangle, payload: V, directory: Directory<V>) -> (),
	_Search: (
		self: RTree<V>,
		rect: Rectangle,
		directory: Directory<V>,
		areaRegistry: { [V]: number }
	) -> (),
	Search: (self: RTree<V>, rect: Rectangle | Rect) -> { V },
	Add: (self: RTree<V>, rect: Rectangle | Rect, payload: V) -> (),
}
export type RTree<V> = typeof(setmetatable(
	{} :: {
		_Master: Directory<V>,
	},
	{} :: RTreeImpl<V>
))

-- Helper Class
-- because the rect class is too slow
local Rectangle = {} :: RectangleImpl
Rectangle.__index = Rectangle

function Rectangle:__tostring(): string
	return self._Key
end

function Rectangle:__eq(other: Rectangle | number): boolean
	if typeof(other) == "number" then
		return self.Area == other
	end
	return self._Key == other._Key
end

function Rectangle:__lt(other: Rectangle | number): boolean
	if typeof(other) == "number" then
		return self.Area < other
	end
	return self._Key < other._Key
end

function Rectangle:__le(other: Rectangle | number): boolean
	if typeof(other) == "number" then
		return self.Area <= other
	end
	return self._Key <= other._Key
end

function Rectangle:__len(): number
	return self.Area
end

function Rectangle:ToRect(other: Rectangle): Rect
	return Rect.new(self.MinX, self.MinY, other.MaxX, other.MaxY)
end

function Rectangle:Contains(other: Rectangle): boolean
	return self.MinX <= other.MinX
		and self.MaxX >= other.MaxX
		and self.MinY <= other.MinY
		and self.MaxY >= other.MaxY
end

function Rectangle:GetOverlap(other: Rectangle): number
	if
		(self.MinX < other.MaxX)
		and (self.MaxX > other.MinX)
		and (self.MinY < other.MaxY)
		and (self.MaxY > other.MinY)
	then
		-- Calculate the overlapping rect's dimensions
		local overlapWidth = math.min(self.MaxX, other.MaxX)
			- math.max(self.MinX, other.MinX)

		local overlapHeight = math.min(self.MaxY, other.MaxY)
			- math.max(self.MinY, other.MinY)

		return overlapWidth * overlapHeight
	else
		return 0
	end
end

function Rectangle:Overlaps(other: Rectangle): boolean
	return (self.MinX < other.MaxX)
		and (self.MaxX > other.MinX)
		and (self.MinY < other.MaxY)
		and (self.MaxY > other.MinY)
end

-- contructors
function newRectangle(minX: number, minY: number, maxX: number, maxY: number): Rectangle
	local height = maxY - minY
	local width = maxX - minX

	local out = setmetatable({
		_Key = `{minX},{minY},{maxX},{maxY}`,
		MinX = minX,
		MinY = minY,
		MaxX = maxX,
		MaxY = maxY,
		Width = width,
		Height = height,
		Area = width * height,
	}, Rectangle)
	table.freeze(out)
	return out
end

function newRectangleFromString(str: string): Rectangle
	local minXStr, minYStr, maxXStr, maxYStr = table.unpack(string.split(str, ","), 1, 4)
	local minX, minY, maxX, maxY =
		tonumber(minXStr), tonumber(minYStr), tonumber(maxXStr), tonumber(maxYStr)
	assert(minX and minY and maxX and maxY, `Invalid rectangle string: {str}`)
	return newRectangle(minX, minY, maxX, maxY)
end

function newRectangleFromRect(rect: Rect): Rectangle
	return newRectangle(rect.Min.X, rect.Min.Y, rect.Max.X, rect.Max.Y)
end

-- Second class
local RTree = {} :: RTreeImpl<unknown>
RTree.__index = RTree

function RTree:_Add(rect: Rectangle, payload: unknown, directory: Directory<unknown>)
	-- add payload to sub-directory
	local subDirectory: Directory<unknown>?
	for subRegionKey, dir in pairs(directory.SubDirectories) do
		if dir.Rectangle:Contains(rect) then
			subDirectory = dir
			break
		end
	end
	-- if sub-directory, add under current directory
	if subDirectory then
		self:_Add(rect, payload, subDirectory)
	else -- if no sub-directory is found, create one
		local freshDirectory = table.freeze({
			Rectangle = rect,
			Item = payload,
			SubDirectories = {},
		}) :: Directory<unknown>

		for r, v in pairs(directory.SubDirectories) do
			if freshDirectory.Rectangle:Contains(v.Rectangle) then
				directory.SubDirectories[r] = nil
				freshDirectory.SubDirectories[r] = v
			end
		end

		directory.SubDirectories[freshDirectory.Rectangle._Key] = freshDirectory
	end
end

function RTree:Add(rect: Rectangle | Rect, payload: unknown)
	local rectangle: Rectangle
	if typeof(rect) == "Rect" then
		rectangle = newRectangleFromRect(rect)
	else
		rectangle = rect
	end
	self:_Add(rectangle, payload, self._Master)
end

function RTree:_Search(
	rect: Rectangle,
	directory: Directory<unknown>,
	areaRegistry: { [unknown]: number }
): ()
	local area = 0
	if directory.Rectangle:Contains(rect) then
		area = directory.Rectangle.Width * directory.Rectangle.Height
	elseif rect:Overlaps(directory.Rectangle) then
		area = rect:GetOverlap(directory.Rectangle)
	end
	if area > 0 then
		if directory.Item ~= nil then
			areaRegistry[directory.Item] = area
		end
		for subReg, subDir in pairs(directory.SubDirectories) do
			self:_Search(rect, subDir, areaRegistry)
		end
	end
end

function RTree:Search(rect: Rectangle | Rect): { unknown }
	local rectangle: Rectangle
	if typeof(rect) == "Rect" then
		rectangle = newRectangleFromRect(rect)
	else
		rectangle = rect
	end
	local areaRegistry: { [unknown]: number } = {}
	self:_Search(rectangle, self._Master, areaRegistry)

	local out: { unknown } = {}
	do
		local i = 0
		for k, v in pairs(areaRegistry) do
			i += 1
			out[i] = k
		end
	end
	table.sort(out, function(a: unknown, b: unknown): boolean
		return areaRegistry[a] < areaRegistry[b]
	end)
	return out
end

return {
	new = function<V>(): RTree<V>
		local self = setmetatable({
			_Master = table.freeze({
				Rectangle = newRectangle(-math.huge, -math.huge, math.huge, math.huge),
				SubDirectories = {},
			}),
		}, RTree :: RTreeImpl<any>)
		table.freeze(self)
		return self
	end,
	Rectangle = {
		new = newRectangle,
		fromString = newRectangleFromString,
		fromRect = newRectangleFromRect,
	},
}