very basic bvh

Bvh.luau

--!strict

-- yet another bvh implementation :P
-- magicoal_nerb

local Queue = require("./Queue")

local Bvh = {}
Bvh.__index = Bvh

export type BvhNode<T> = {
	min: vector,
	max: vector,
	value: T?,
	
	left: number,
	right: number,
}

export type Bvh<T> = typeof(setmetatable({} :: {
	nodes: { BvhNode<T> },
	size: number,
	queue: Queue.Queue<T>,
}, Bvh))

local function doesAABBCollideAABB(
	min0: vector,
	max0: vector,
	min1: vector,
	max1: vector
): boolean
	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 getGreatestAxis(axis: vector): vector
	local x = math.abs(axis.x)
	local y = math.abs(axis.y)
	local z = math.abs(axis.z)
	
	if x > y and x > z then
		return vector.create(math.sign(axis.x), 0, 0)
	elseif y > x and y > z then
		return vector.create(0, math.sign(axis.y), 0)
	else
		return vector.create(0, 0, math.sign(axis.z))
	end
end

local function makePair(id: number, a: number, b: number): number
	-- encodes the pair as a number
	assert(a < 256 and b < 256)
	return bit32.lshift(id, 16)
		+ bit32.lshift(a, 8)
		+ b
end

local function getPair(pair: number): (number, number, number)
	-- decodes the number pair for us !!
	return bit32.rshift(pair, 16),
		bit32.band(bit32.rshift(pair, 8), 0xFF),
		bit32.band(pair, 0xFF)
end

function Bvh.new<T>(): Bvh<T>
	return setmetatable({ 
		nodes = {},
		queue = Queue.new(),
		size = 0,
	}, Bvh)
end

function Bvh.build<T>(self: Bvh<T>, nodes: { BvhNode<T> })
	local buildQueue = Queue.new() :: Queue.Queue<number>
	buildQueue:enqueue(makePair(1, 1, #nodes))
	
	-- the build process uses a very simple heuristic,
	-- we will split based upon the center of the greatest axis
	-- in the subset. no sah used, but this is fine for small scenes
	-- and it is much faster than a full SAH split
	
	local output = self.nodes
	while not buildQueue:empty() do
		local pair = buildQueue:dequeue()
		local index, leftIndex, rightIndex = getPair(pair)
		if rightIndex - leftIndex == 1 then
			-- we made a leaf with 2 elements :P
			local leftNode = nodes[leftIndex]
			local rightNode = nodes[rightIndex]
			local min = vector.min(leftNode.min, rightNode.min)
			local max = vector.max(leftNode.max, rightNode.max)
			output[2*index + 2] = rightNode
			output[2*index + 1] = leftNode
			
			-- create the branch node here!!
			output[index] = {
				min = min,
				max = max,
				left = 2*index + 1,
				right = 2*index + 2,
				value = nil,
			}
						
			continue
		elseif leftIndex == rightIndex then
			-- otherwise, we can just set it to the current
			-- node provided to us
			output[index] = nodes[leftIndex]
			continue
		end

		-- expand the bounding boxes
		local min = nodes[leftIndex].min
		local max = nodes[leftIndex].max
		for i = leftIndex + 1, rightIndex do
			local node = nodes[i]
			min = vector.min(min, node.min)
			max = vector.max(max, node.max)
		end
		
		local delta = max - min
		local center = (min + max) * 0.5
		local axis = getGreatestAxis(delta)
		
		-- partition stuff on the left side !!
		local left = leftIndex
		local temp = {}
		for i = leftIndex, rightIndex do
			-- do the center split
			local node = nodes[i]
			local centroid = (node.min + node.max) * 0.5
			if vector.dot(centroid - center, axis) < 0 then
				nodes[left] = node
				left += 1
			else
				table.insert(temp, node)
			end
		end
		
		-- then we add back the right elements
		local t = {}
		for i = left, rightIndex do
			nodes[i] = temp[i - left + 1]
		end
		
		for i = leftIndex, rightIndex do
			if t[nodes[i]] then
				error(i - left)
			end
			t[nodes[i]] = true
		end
		
		-- add branch node here
		output[index] = {
			min = min,
			max = max,
			left = 2*index + 1,
			right = 2*index + 2,
			value = nil,
		}
		
		if left > rightIndex then
			-- just get the center if we're above the right index
			-- because it's better than degenerating into a linked list
			left = (leftIndex + rightIndex) // 2
		end
		
		buildQueue:enqueue(makePair(2*index + 1, leftIndex, left - 1))
		buildQueue:enqueue(makePair(2*index + 2, left, rightIndex))
	end
end

function Bvh.query<T>(self: Bvh<T>, min: vector, max: vector): { T }
	local queue = self.queue
	local nodes = self.nodes
	
	queue:clear()
	queue:enqueue(1)
	
	local output = {} :: { T }
	while not queue:empty() do
		local id = queue:dequeue()
		if id == 0 then
			continue
		end
		
		-- get the node from here instead
		local node = nodes[id]
		if node.value then
			-- we reached a leaf node, so add
			-- to the output
			table.insert(output, node.value)
			continue
		end
		
		local rightNode = nodes[node.right]
		local leftNode = nodes[node.left]
		if leftNode and doesAABBCollideAABB(min, max, leftNode.min, leftNode.max) then
			-- we are colliding with the left node(if it exists)
			queue:enqueue(node.left)
		end
		
		if rightNode and doesAABBCollideAABB(min, max, rightNode.min, rightNode.max) then
			-- we are colliding with the right node(if it exists)
			queue:enqueue(node.right)
		end
	end
	
	return output
end

return Bvh

Queue.luau

--!strict

-- yet another queue implementation :P
-- magicoal_nerb

local Queue = {}
Queue.__index = Queue

export type Queue<T> = typeof(setmetatable({} :: {
	data: { T },
	size: number,
}, Queue))

function Queue.new<T>(): Queue<T>
	return setmetatable({
		data = {},
		size = 0,
	}, Queue)
end

function Queue.empty<T>(self: Queue<T>): boolean
	return self.size == 0
end

function Queue.iterate<T>(self: Queue<T>, callback: (T, ...any) -> (), ...: any)
	for i, object in self.data do
		callback(object, ...)
	end
end

function Queue.enqueue<T>(self: Queue<T>, data: T)
	self.data[self.size + 1] = data
	self.size += 1
end

function Queue.dequeue<T>(self: Queue<T>): T
	assert(self.size > 0, "attempted to dequeue empty queue lol")
	
	self.size -= 1
	return self.data[self.size + 1]
end

function Queue.clear<T>(self: Queue<T>)
	self.size = 0
	table.clear(self.data)
end

return Queue