Files
rust-raytracer/src/bvh.rs

443 lines
16 KiB
Rust

use crate::{node::Node, ray::*, EPSILON};
use nalgebra::{distance, Matrix4, Point3, Vector3};
use std::collections::HashMap;
use std::fmt;
// Debuging statics
static mut STATIC0: i32 = 0;
static mut STATIC1: i32 = 0;
static mut STATIC2: i32 = 0;
static mut STATIC3: i32 = 0;
static mut STATIC4: i32 = 0;
// BOUNDING BOX -----------------------------------------------------------------
#[derive(Clone)]
pub struct AABB {
pub bln: Point3<f64>,
pub trf: Point3<f64>,
pub centroid: Point3<f64>,
}
impl AABB {
// New box with respective coordinates
pub fn new(bln: Point3<f64>, trf: Point3<f64>) -> AABB {
let bln = bln - Vector3::new(EPSILON, EPSILON, EPSILON);
let trf = trf + Vector3::new(EPSILON, EPSILON, EPSILON);
let centroid = bln + (trf - bln) / 2.0;
AABB { bln, trf, centroid }
}
//Empty box
pub fn empty() -> AABB {
AABB {
bln: Point3::new(f64::MAX, f64::MAX, f64::MAX),
trf: Point3::new(f64::MIN, f64::MIN, f64::MIN),
centroid: Point3::new(0.0, 0.0, 0.0),
}
}
//Apply a matrix transformation to a box
pub fn transform_mut(&mut self, mat: &Matrix4<f64>) {
let bln = &mut self.bln;
let trf = &mut self.trf;
let centroid = &mut self.centroid;
self.bln = mat.transform_point(bln);
self.trf = mat.transform_point(trf);
self.centroid = mat.transform_point(centroid);
}
// Intersect bounding box exactly
pub fn intersect_ray(&self, ray: &Ray) -> bool {
let bln = &self.bln;
let trf = &self.trf;
let t1 = (bln - ray.a).component_div(&ray.b);
let t2 = (trf - ray.a).component_div(&ray.b);
let tmin = t1.inf(&t2).max();
let tmax = t1.sup(&t2).min();
tmax >= tmin && tmax > 0.0
}
// Intersect ray with some epsilon tolerance
pub fn intersect_ray_aprox(&self, ray: &Ray) -> bool {
let bln = &self.bln;
let trf = &self.trf;
let t1 = (bln - ray.a).component_div(&ray.b);
let t2 = (trf - ray.a).component_div(&ray.b);
let tmin = t1.inf(&t2).max();
let tmax = t1.sup(&t2).min();
tmax >= tmin - EPSILON && tmax > -EPSILON
}
// Get the center of this bounding box
fn get_centroid(&self) -> Point3<f64> {
self.centroid
}
// Make a new AABB that contains both
pub fn join(&self, other: &AABB) -> AABB {
AABB::new(
Point3::new(
self.bln.x.min(other.bln.x),
self.bln.y.min(other.bln.y),
self.bln.z.min(other.bln.z),
),
Point3::new(
self.trf.x.max(other.trf.x),
self.trf.y.max(other.trf.y),
self.trf.z.max(other.trf.z),
),
)
}
//Join mutably
pub fn join_mut(&mut self, other: &AABB) {
self.bln = Point3::new(
self.bln.x.min(other.bln.x),
self.bln.y.min(other.bln.y),
self.bln.z.min(other.bln.z),
);
self.trf = Point3::new(
self.trf.x.max(other.trf.x),
self.trf.y.max(other.trf.y),
self.trf.z.max(other.trf.z),
);
self.centroid = self.bln + (self.trf - self.bln) / 2.0;
}
//Grow the AABB to contain the cover the point
pub fn grow(&self, other: &Point3<f64>) -> AABB {
AABB::new(
Point3::new(
self.bln.x.min(other.x),
self.bln.y.min(other.y),
self.bln.z.min(other.z),
),
Point3::new(
self.trf.x.max(other.x),
self.trf.y.max(other.y),
self.trf.z.max(other.z),
),
)
}
//Grow mutably
pub fn grow_mut(&mut self, other: &Point3<f64>) {
self.bln = Point3::new(
self.bln.x.min(other.x),
self.bln.y.min(other.y),
self.bln.z.min(other.z),
);
self.trf = Point3::new(
self.trf.x.max(other.x),
self.trf.y.max(other.y),
self.trf.z.max(other.z),
);
self.centroid = self.bln + (self.trf - self.bln) / 2.0;
}
// Size of AABB
pub fn size(&self) -> Vector3<f64> {
self.trf - self.bln
} //Surface area of AABB
pub fn surface_area(&self) -> f64 {
let size = self.size();
2.0 * (size.x * size.y + size.x * size.z + size.y * size.z)
}
pub fn area(&self) -> f64 {
let extent = self.trf - self.bln;
return extent.x * extent.y + extent.y * extent.z + extent.z * extent.x;
}
// Volume of the AABB
pub fn volume(&self) -> f64 {
let size = self.size();
size.x * size.y * size.z
}
}
impl fmt::Display for AABB {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
if self.bln[0] == f64::MAX || self.trf[0] == f64::MIN {
writeln!(f, "Empty aabb")
} else {
writeln!(f, "bln: {}\ntrf: {}", self.bln, self.trf)
}
}
}
#[derive(Clone)]
pub struct BVHNode {
aabb: AABB, //The nodes bounding box
l_idx: usize, //Child node l, the right node is alway l_idx + 1
first_prim: usize, //First primitive that the node encapsulates
prim_count: usize, //Number of primitives the node encapsulates
}
impl BVHNode {
pub fn default() -> BVHNode {
BVHNode {
aabb: AABB::empty(),
l_idx: 0,
first_prim: 0,
prim_count: 0,
}
}
}
impl fmt::Display for BVHNode {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
writeln!(f, "l_idx: {}", self.l_idx)?;
writeln!(f, "First Prim: {}", self.first_prim)?;
writeln!(f, "Prim Count: {}", self.prim_count)?;
writeln!(f, "aabb: {}", self.aabb)
}
}
pub struct BVH {
bvh_nodes: Vec<BVHNode>, //BVH nodes with AABBs
nodes: Vec<Node>, //Nodes with primitives
nodes_used: usize,
}
impl BVH {
//Build a bvh by subdividing recursively
pub fn build(in_nodes: &HashMap<String, Node>) -> BVH {
/*
Make our own vec of nodes so that we can refer to it by index
This might be expensive so another method is preferred
*/
let mut nodes = vec![];
for (_, node) in in_nodes {
nodes.push(node.clone());
}
//A BVH tree will be maximum size of 2*n + 1
//Initialise an empty BVHNode with empty AABB
let n = nodes.len();
let bvh_nodes: Vec<BVHNode> = vec![BVHNode::default(); 2 * n + 1];
//Begin constructing our BVH tree
//One node used to begin with (The root node)
let nodes_used = 1;
let mut tree = BVH {
nodes,
bvh_nodes,
nodes_used,
};
// Get the root node at index 0
let root = &mut tree.bvh_nodes[0];
root.l_idx = 0; //Root node has no left or right child to begin
(root.first_prim, root.prim_count) = (0, n); //Make root include all n nodes
tree.update_bvh_node_aabb(0); //Create the root nodes AABB on the n primitives
tree.subdivide(0); //Sub divide the root node
tree
}
// Will update the node's AABB at bvh_nodes[index]
fn update_bvh_node_aabb(&mut self, index: usize) {
// We will make his node bound all its primitives
let bvh_node = &mut self.bvh_nodes[index]; // Current BVHNode
let bvh_node_aabb = &mut bvh_node.aabb; //Current node AABB
let first_prim = bvh_node.first_prim; //Start index of prim
let prim_count = bvh_node.prim_count; //Number of primitives within the nodes aabb
for i in 0..prim_count {
let node = &self.nodes[first_prim + i]; //Get the node from the Vec<Node>
bvh_node_aabb.join_mut(&node.aabb); //Join it with the BVH node's AABB
}
// unsafe {
// println!("UPDATE TO AABB ---- {STATIC0}");
// STATIC0 += 1;
// let bvh_node = &mut self.bvh_nodes[index]; //Get the BVHNode we are working on
// println!("{bvh_node}");
// }
}
// Subdivision, will subdivide a split
fn subdivide(&mut self, index: usize) {
//Get the bvh_node we will be altering
// Determine the axis and position of the split plane
// Split the group of primitives in two halves using the split plane
// Create child nodes for each half
// Recurse into each of the child nodes.
//Leaf node case, we cannot sub-divide any more
if self.bvh_nodes[index].prim_count == 1 {
return;
};
/* ------------ SUBDIVIDE BY LONGEST AXIS ------------ */
//Get information about the node we want to subdivide
let (bln, trf) = (
self.bvh_nodes[index].aabb.bln,
self.bvh_nodes[index].aabb.trf,
);
let extent = trf - bln;
let mut axis = 0; // Assume that x is longest
if extent.y > extent.x {
axis = 1; // Split y if longest
};
if extent.z > extent[axis] {
axis = 2; // Split z if longest
};
let split_pos = bln[axis] + extent[axis] * 0.5; // Final split down the middle of AABB
/* --------- SUBDIVIDE BY Surface Area Heuristic ---------*/
// let mut best_axis: Option<usize> = None;
// let mut best_pos = 0.0;
// let mut best_cost = 1e30;
// let first_prim_idx = self.bvh_nodes[index].first_prim;
// for axis in 0..3 {
// for i in 0..self.bvh_nodes[index].prim_count {
// let node = &self.nodes[first_prim_idx + i];
// //Get the centroid of the bounding box
// let centroid = node.aabb.get_centroid();
// //Get the candidate position
// let candidate_pos = world_centroid[axis];
// let cost = self.evaluate_sah(&self.bvh_nodes[index], axis, candidate_pos);
// if cost < best_cost {
// best_pos = candidate_pos;
// best_axis = Some(axis);
// best_cost = cost;
// }
// }
// }
// let axis = match best_axis {
// Some(axis) => axis,
// None => 0,
// };
// let split_pos = best_pos;
let left_count;
let right_count;
let mut i;
let mut j;
{
let bvh_node = &mut self.bvh_nodes[index];
i = bvh_node.first_prim; //Start of array
j = i + bvh_node.prim_count - 1; //End of array
while i <= j {
//Perform a quicksort dependent on location
let node = &self.nodes[i]; // Node we would like to sort
let centroid = node.aabb.get_centroid(); //Centroid of node we would like to sort
if centroid[axis] < split_pos {
i += 1; // On Left-Hand-Side
} else {
self.nodes.swap(i, j);
j -= 1; // On Right-Hand-Side
}
}
//Now we have two splits
//The lhs of the array is in the left split 0..left_count
//The rhs of the array is on the right split left_count + 1..n
left_count = i - bvh_node.first_prim; //Number of prims on lhs
right_count = bvh_node.prim_count - left_count;
//println!("SPLIT INTO: {left_count} {right_count}");
if left_count == 0 || left_count == bvh_node.prim_count {
//Split did nothing
return;
}
}
// unsafe {
// println!("SUBDIVIDE: {STATIC1}");
// println!("SPLIT INTO: {left_count} ");
// STATIC1 += 1;
// }
let l_idx = self.nodes_used; //Left child
self.bvh_nodes[index].l_idx = l_idx;
self.nodes_used = self.nodes_used + 2;
//Set left node information
self.bvh_nodes[l_idx].first_prim = self.bvh_nodes[index].first_prim; //Left split begins at parent split
self.bvh_nodes[l_idx].prim_count = left_count; // Left prims
//Set right node information
self.bvh_nodes[l_idx + 1].first_prim = i; // Right split start index
self.bvh_nodes[l_idx + 1].prim_count = right_count;
//Current node is not a leaf node
self.bvh_nodes[index].prim_count = 0;
self.update_bvh_node_aabb(l_idx); //Update AABB for left of split
self.update_bvh_node_aabb(l_idx + 1); //Update AABB for right of split
//Recurse
self.subdivide(l_idx); // Subdivide left index
self.subdivide(l_idx + 1); // SUbdivide right index
}
// Traverse the BVH, 0 will be needed to start at root node
pub fn traverse(&self, ray: &Ray, idx: usize) -> Option<(&Node, Intersection)> {
let bvh_node = &self.bvh_nodes[idx];
if !bvh_node.aabb.intersect_ray(&ray) {
// No intersection with BVH in world coordinates
return None;
}
if bvh_node.prim_count != 0 {
// Leaf node intersection — test all primitives in the leaf
let mut closest: Option<(&Node, Intersection)> = None;
let mut closest_dist = f64::MAX;
for i in 0..bvh_node.prim_count {
let node = &self.nodes[bvh_node.first_prim + i];
if !node.active {
continue;
}
if let Some(intersect) = node.intersect_ray(&ray) {
if intersect.distance < EPSILON {
continue;
}
if intersect.distance < closest_dist {
closest_dist = intersect.distance;
closest = Some((node, intersect));
}
}
}
return closest;
} else {
//Recurse down the BVH
//Recurse down the BVH right node
let intersect_l = self.traverse(ray, bvh_node.l_idx);
let intersect_r = self.traverse(ray, bvh_node.l_idx + 1);
match (intersect_l, intersect_r) {
(None, None) => return None,
(Some(intersect), None) => return Some(intersect),
(None, Some(intersect)) => return Some(intersect),
(Some((node_l, inter_l)), Some((node_r, inter_r))) => {
//Compare intersect distance
let dist_l = distance(&ray.a, &inter_l.point);
let dist_r = distance(&ray.a, &inter_r.point);
if dist_l < dist_r {
return Some((node_l, inter_l));
} else {
return Some((node_r, inter_r));
}
}
}
}
}
fn evaluate_sah(&self, node: &BVHNode, axis: usize, pos: f64) -> f64 {
// determine triangle counts and bounds for this split candidate
let mut l_aabb = AABB::empty();
let mut r_aabb = AABB::empty();
let mut l_count = 0;
let mut r_count = 0;
for i in 0..node.prim_count {
let aabb = self.nodes[node.first_prim + i].get_world_aabb();
if aabb.trf[axis] < pos {
l_count += 1;
l_aabb.join_mut(&aabb);
} else {
r_count += 1;
r_aabb.join_mut(&aabb);
}
}
let cost = l_count as f64 * l_aabb.area() + r_count as f64 * r_aabb.area();
match cost > 0.0 {
true => cost,
false => 1e30,
}
}
}
impl fmt::Display for BVH {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
for (i, node) in self.bvh_nodes.iter().enumerate() {
writeln!(f, "Node: {i}")?;
writeln!(f, "{node}")?;
}
write!(f, "")
}
}