Octree.lua

--!strict


-- // Helper Functions

-- Insert Data & Point into Node
local function InsertHelper(self: Octree, Node: Node, Point: Vector3, Data: any, Depth: number): ()
	while true do
		-- Insert the Data directly into the Node
		if Depth == self.MaxDepth or #Node.DataTables >= self.MaxObjectsPerNode then
			-- Check if the Point already exists in the Node's DataTables
			local PointExists: boolean = false
			for _, DataTable: DataTable in ipairs(Node.DataTables) do
				if DataTable.Point == Point then
					DataTable.Data = Data
					PointExists = true
					break
				end
			end
			
			-- If the Point was not found in the Node's DataTables, then insert it
			if not PointExists then
				table.insert(Node.DataTables, {Point = Point, Data = Data} :: DataTable)
			end
			
			break
		end
		
		
		-- Divide the Node into Octants and Insert the Data into the Appropriate Octant
		local Octant: number = self:GetOctant(Node.Position, Node.Size, Point)
		
		if not Node.Children[Octant] then
			Node.Children[Octant] = Node.new(self:GetOctantBounds(Node.Position, Node.Size, Octant))
		end
		
		
		-- Recurse for Children
		Node = Node.Children[Octant]
		Depth = Depth + 1
	end
end

-- Remove Node with Data
local function RemoveHelper(self: Octree, Node: Node, Data: any): Node
	for Index: number, DataTable: DataTable in ipairs(Node.DataTables) do
		if DataTable.Data == Data then
			table.remove(Node.DataTables, Index)
			break
		end
	end
	
	
	local NumChildren: number = #Node.Children
	
	if NumChildren == 0 then
		return Node
	end
	
	
	for Index: number, Child: Node in ipairs(Node.Children) do
		Node.Children[Index] = RemoveHelper(self, Child, Data)
	end
	
	-- If the Node has No Data and All its Children have no Data, then Prune the Node
	if #Node.DataTables == 0 and NumChildren == #Node.Children then
		return Node
	end
	
	
	return Node
end

-- Remove Node with Point
local function RemoveByPointHelper(self: Octree, Node: Node, Point: Vector3): Node
	for Index: number, DataTable: DataTable in ipairs(Node.DataTables) do
		if DataTable.Point == Point then
			table.remove(Node.DataTables, Index)
			break
		end
	end
	
	local NumChildren: number = #Node.Children
	
	if NumChildren == 0 then
		return Node
	end
	
	
	for Index: number, Child: Node in ipairs(Node.Children) do
		Node.Children[Index] = RemoveByPointHelper(self, Child, Point)
	end
	
	-- If the Node has No Data and All its Children have no Data, then Prune the Node
	if #Node.DataTables == 0 and NumChildren == #Node.Children then
		return Node
	end
	
	
	return Node
end

-- Recursively Traverse the Tree and Return the Data for all Nodes that Intersect with the Query Region
local function QueryHelper(self: Octree, Node: Node, Position: Vector3, Size: Vector3, Results: {DataTable}): ()
	if self:Intersects(Node.Position, Node.Size, Position, Size) then
		-- Add the Data from this Node to the Results
		for _, DataTable: DataTable in ipairs(Node.DataTables) do
			table.insert(Results, DataTable)
		end
		
		-- Recurse for Children
		for _, Child: Node in ipairs(Node.Children) do
			QueryHelper(self, Child, Position, Size, Results)
		end
	end
end


-- // Node Class

local Node = {}
Node.__index = Node

-- Construct new Node Object
function Node.new(Position: Vector3, Size: Vector3): Node
	return setmetatable({
		Position = Position,
		Size = Size,
		Children = {},
		DataTables = {}
	}, Node)
end


-- // Octree Class

local Octree = {}
Octree.__index = Octree

-- Construct new Octree Object
function Octree.new(Position: Vector3, Size: Vector3, MaxDepth: number, MaxObjects: number): Octree
	return setmetatable({
		Root = Node.new(Position, Size),
		MaxDepth = MaxDepth,
		MaxObjects = MaxObjects
	}, Octree)
end

-- Insert Data for Corresponding Octant in Tree
function Octree.Insert(self: Octree, Point: Vector3, Data: any): ()
	InsertHelper(self, self.Root, Point, Data, 0)
end

-- Remove Octant for Corresponding Data in Tree
function Octree.Remove(self: Octree, Data: any): ()
	self.Root = RemoveHelper(self, self.Root, Data)
end

-- Remove Octant for Corresponding Point in Tree
function Octree.RemoveByPoint(self: Octree, Point: Vector3): ()
	self.Root = RemoveByPointHelper(self, self.Root, Point)
end

-- Resets Root, Clearing All Nodes
function Octree.Clear(self: Octree): ()
	self.Root = Node.new(self.Root.Position, self.Root.Size)
end

-- Return Array of DataTable for Octants in Bounds
function Octree.Query(self: Octree, Position: Vector3, Size: Vector3): {DataTable}
	local Results: {DataTable} = table.create(#self.Root)
	
	QueryHelper(self, self.Root, Position, Size, Results)
	
	return Results
end

-- Find the Nearest Data Point to the Given Point within the Given Maximum Distance
function Octree.GetNearestDataPoint(self: Octree, Point: Vector3, MaxDistance: number): (Vector3, any)
	-- Query the Tree for Data within the Maximum Distance
	local Results: {DataTable} = self:Query(Point, Vector3.new(MaxDistance, MaxDistance, MaxDistance))
	
	-- Linear Search for the Nearest Data Point
	local NearestPoint: Vector3 = nil
	local NearestData: any = nil
	local MinDistance: number = MaxDistance
	
	for _, DataTable: DataTable in ipairs(Results) do
		local Distance: number = (DataTable.Point - Point).Magnitude
		
		if Distance < MinDistance then
			MinDistance = Distance
			NearestPoint = DataTable.Point
			NearestData = DataTable.Data
		end
	end
	
	
	return NearestPoint, NearestData
end

-- Determine which Octant a Point Lies In
function Octree.GetOctant(self: Octree, Position: Vector3, Size: Vector3, Point: Vector3): number
	local CX: number, CY: number, CZ: number = Position.X, Position.Y, Position.Z
	local PX: number, PY: number, PZ: number = Point.X, Point.Y, Point.Z
	
	local PX_LESS_CX: boolean = PX < CX
	local PY_LESS_CY: boolean = PY < CY
	local PZ_LESS_CZ: boolean = PZ < CZ
	
	
	if PX_LESS_CX and PY_LESS_CY and PZ_LESS_CZ then
		return 1
	elseif (not PX_LESS_CX) and PY_LESS_CY and PZ_LESS_CZ then
		return 2
	elseif PX_LESS_CX and (not PY_LESS_CY) and PZ_LESS_CZ then
		return 3
	elseif (not PX_LESS_CX) and (not PY_LESS_CY) and PZ_LESS_CZ then
		return 4
	elseif PX_LESS_CX and PY_LESS_CY and not (PZ_LESS_CZ) then
		return 5
	elseif (not PX_LESS_CX) and PY_LESS_CY and (not PZ_LESS_CZ) then
		return 6
	elseif PX_LESS_CX and (not PY_LESS_CY) and (not PZ_LESS_CZ) then
		return 7
	end
	
	
	return 8
end

-- Helper function to get the bounding box for a given octant
-- Returns Position & Size of Given Octant
function Octree.GetOctantBounds(self: Octree, Position: Vector3, Size: Vector3, Octant: number): (Vector3, Vector3)
	local CX: number, CY: number, CZ: number = Position.X, Position.Y, Position.Z
	local WX: number, WY: number, WZ: number = Size.X, Size.Y, Size.Z
	
	local WX_HALF: number = WX * 0.5
	local WY_HALF: number = WY * 0.5
	local WZ_HALF: number = WZ * 0.5
	
	local CX_WX_HALF: number = CX - WX_HALF
	local CY_WY_HALF: number = CY - WY_HALF
	local CZ_WZ_HALF: number = CZ - WZ_HALF
	
	
	local ReturnPosition: Vector3
	local ReturnSize: Vector3 = Vector3.new(WX_HALF, WY_HALF, WZ_HALF)
	
	
	if Octant == 1 then
		ReturnPosition = Vector3.new(CX_WX_HALF, CY_WY_HALF, CZ_WZ_HALF)
	end
	
	if Octant == 2 then	
		ReturnPosition = Vector3.new(CX, CY_WY_HALF, CZ_WZ_HALF)
	end
	
	if Octant == 3 then
		ReturnPosition = Vector3.new(CX_WX_HALF, CY, CZ_WZ_HALF)
	end
	
	if Octant == 4 then
		ReturnPosition = Vector3.new(CX, CY, CZ_WZ_HALF)
	end
	
	if Octant == 5 then
		ReturnPosition = Vector3.new(CX_WX_HALF, CY_WY_HALF, CZ)
	end
	
	if Octant == 6 then
		ReturnPosition = Vector3.new(CX, CY_WY_HALF, CZ)
	end
	
	if Octant == 7 then
		ReturnPosition = Vector3.new(CX_WX_HALF, CY, CZ)
	end
	
	ReturnPosition = Vector3.new(CX, CY, CZ)
	
	
	return ReturnPosition, ReturnSize
end

-- Checks if Two Bounds Intersect
function Octree.Intersects(self: Octree, P1: Vector3, S1: Vector3, P2: Vector3, S2: Vector3): boolean
	local X1: number, Y1: number, Z1: number = P1.X, P1.Y, P1.Z
	local W1: number, H1: number, D1: number = S1.X, S1.Y, S1.Z
	
	local X2: number, Y2: number, Z2: number = P2.X, P2.Y, P2.Z
	local W2: number, H2: number, D2: number = S2.X, S2.Y, S2.Z
	
	-- Intersects along X-Axis
	if X1 + W1 < X2 or X2 + W2 < X1 then
		return false
	end
	
	-- Intersects along Y-Axis
	if Y1 + H1 < Y2 or Y2 + H2 < Y1 then
		return false
	end
	
	-- Intersects along Z-Axis
	if Z1 + D1 < Z2 or Z2 + D2 < Z1 then
		return false
	end
	
	
	return true
end


-- // Types

type DataTable = {Point: Vector3, Data: any}

export type Node = typeof(setmetatable({} :: {
	Position: Vector3,
	Size: Vector3,
	Children: {Node},
	DataTables: {DataTable}
}, Node))

export type Octree = typeof(setmetatable({} :: {
	Root: Node,
	MaxDepth: number,
	MaxObjects: number
}, Octree))


return Octree