morton-octree.luau

--!strict
local Queue = require("./Queue")

local Octree = {}
Octree.__index = Octree

export type OctreeNode = {
	-- nodes expect 8 children if childIndex != 0
	min: vector,
	max: vector,
	childIndex: number, 
}

export type OctreePair = {
	left: number,
	right: number,
	depth: number,
	id: number,
}

export type Wrapped<T> = T & OctreeNode
export type Octree<T> = typeof(setmetatable({} :: {
	queue: Queue.Queue<number>,
	data: { OctreeNode | Wrapped<T> },
}, Octree))

local function doesAABBCollideAABB(
	aabb0: OctreeNode,
	min1: vector,
	max1: vector
): boolean
	-- Checks if two aabbs collide
	local min0 = aabb0.min
	local max0 = aabb0.max
	return min0.x <= max1.x and max0.x >= min1.x
		and min0.y <= max1.y and max0.y >= min1.y
		and min0.z <= max1.z and max0.z >= min1.z
end

local function expandBits(v: number): number
	-- Expands a 10 bit number into triples
	v = bit32.band(v * 0x00010001, 0xFF0000FF)
	v = bit32.band(v * 0x00000101, 0x0F00F00F)
	v = bit32.band(v * 0x00000011, 0xC30C30C3)
	v = bit32.band(v * 0x00000005, 0x49249249)
	return v
end

local function mortonCode(
	x: number,
	y: number,
	z: number
): number
	-- Expands three, 10 bit numbers into a 30 bit number
	-- Assumes x,y,z are from 0 to 1023 (inclusive)
	return expandBits(x) * 4
		+ expandBits(y) * 2
		+ expandBits(z)
end

local function encodeNodes<T>(nodes: { Wrapped<T> })
	-- Encodes our nodes within a morton code
	local max = -vector.one * math.huge
	local min = vector.one * math.huge
	for i, node in nodes do
		local position = (node.min + node.max) * 0.5
		min = vector.min(min, position)
		max = vector.max(max, position)
	end

	local range = (vector.one * 1023) / (max - min)
	for i, node in nodes do
		local center = (node.min + node.max) * 0.5
		local position = vector.floor((center - min) * range)
		node.childIndex = mortonCode(
			position.x,
			position.y,
			position.z
		)
	end
end

function Octree.new<T>(
	nodes: { Wrapped<T> },
	radius: number
): Octree<T>
	return setmetatable({
		queue = Queue.new(8, 0),
		data = {},
	}, Octree):build(nodes)
end

function Octree.query<T>(
	self: Octree<T>,
	min: vector,
	max: vector,
	callback: (node: Wrapped<T>) -> ()
)
	local data = self.data
	local q = self.queue
	q:clear()
	q:enqueue(1)
	
	while not q:empty() do
		-- Get current node
		local node = data[q:dequeue()]
		local childIndex = node.childIndex - 1
		for i = 1, 8 do
			-- Loop through the node's children, check
			-- if it collides with our region
			local id = childIndex + i
			local child = data[id]
			if child and doesAABBCollideAABB(child, min, max) then
				-- Collided with our region!
				if child.childIndex == 0 then
					callback(child :: Wrapped<T>)
				else
					q:enqueue(id)
				end
			end
		end
	end
end

function Octree.build<T>(
	self: Octree<T>,
	nodes: { Wrapped<T> }
): Octree<T>
	-- Create the build queue
	local q = Queue.new(8, {} :: OctreePair)
	q:enqueue({
		left = 1,
		right = #nodes,
		
		depth = 1,
		id = 1,
	})
	
	-- Create arrays
	local temp = table.create(#nodes)
	local sums = table.create(8, 0)
	local bins = table.create(8, 0)
	encodeNodes(nodes)
	
	local count = 1
	while not q:empty() do
		local pair = q:dequeue()
		local depth = 30 - 3*pair.depth
		local mask = bit32.lshift(0x7, depth)
		assert(depth > 0, "max depth limit reached")
		
		-- Get the bounding box of the branch
		local max = -vector.one * math.huge
		local min = vector.one * math.huge
		for i = pair.left, pair.right do
			local node = nodes[i]
			local code = bit32.band(node.childIndex, mask)
			
			bins[bit32.rshift(code, depth) + 1] += 1
			min = vector.min(min, node.min)
			max = vector.max(max, node.max)
		end
		
		-- Add data
		local childIndex = count + 1
		self.data[pair.id] = {
			childIndex = childIndex,
			min = min,
			max = max,
		}

		-- For each bin, check if they need to be
		-- subdivided more.		
		local sum = pair.left
		for i = 1, 8 do
			sums[i] = sum
			sum += bins[i]
		end
		
		-- Partition nodes
		for i = pair.left, pair.right do
			local node = nodes[i]
			
			local code = bit32.band(node.childIndex, mask)
			local id = bit32.rshift(code, depth) + 1
			temp[sums[id]] = node
			sums[id] += 1
		end
		
		for i = pair.left, pair.right do
			nodes[i] = temp[i]
		end
		
		-- Finalize the bins
		local left = pair.left
		count += 8
		
		for i = 1, 8 do
			local size = bins[i]
			local right = left + size - 1
			if size == 0 then
				-- Nothing... I suppose
			elseif size == 1 then
				local leaf = nodes[left]
				leaf.childIndex = 0
				
				self.data[childIndex + i - 1] = leaf
			elseif size <= 8 then
				-- Leaves
				local leafMax = -vector.one * math.huge
				local leafMin = vector.one * math.huge
				local leafIndex = count + 1
				
				for j = left, right do
					local leaf = nodes[j]
					leaf.childIndex = 0

					leafMin = vector.min(leafMin, leaf.min)
					leafMax = vector.max(leafMax, leaf.max)
					self.data[leafIndex + j - left] = leaf
				end
				
				self.data[childIndex + i - 1] = {
					min = leafMin,
					max = leafMax,
					childIndex = leafIndex
				}
				
				count += 8
			else
				-- Branch
				q:enqueue({
					left = left,
					right = right,
					
					depth = pair.depth + 1,
					id = childIndex + i - 1
				})
			end
			
			-- Zero out the bin
			left = right + 1
			bins[i] = 0
		end
	end
	
	return self
end

return Octree