From 18bb6a250bbe704df3543d1c74045f3bea23e3ea Mon Sep 17 00:00:00 2001 From: Ishan Jain Date: Tue, 2 Mar 2021 01:25:27 +0530 Subject: [PATCH] Completed AABB BVH implementation --- src/aabb.rs | 51 ++++++++++++++ src/bvh.rs | 134 +++++++++++++++++++++++++++++++++++++ src/demo.rs | 10 +-- src/hitable.rs | 20 +++++- src/hitable_list.rs | 24 ++++++- src/main.rs | 10 ++- src/types/moving_sphere.rs | 10 ++- src/types/sphere.rs | 7 +- src/types/vec3.rs | 2 +- 9 files changed, 253 insertions(+), 15 deletions(-) create mode 100644 src/aabb.rs create mode 100644 src/bvh.rs diff --git a/src/aabb.rs b/src/aabb.rs new file mode 100644 index 0000000..15b9e15 --- /dev/null +++ b/src/aabb.rs @@ -0,0 +1,51 @@ +use crate::types::{Ray, Vec3}; + +#[derive(Debug, Copy, Clone)] +pub struct Aabb { + pub min: Vec3, + pub max: Vec3, +} + +impl Aabb { + pub const fn new(min: Vec3, max: Vec3) -> Self { + Self { min, max } + } + + pub fn hit(&self, ray: &Ray, mut t_min: f64, mut t_max: f64) -> bool { + for i in 0..=2 { + let inverse_dir = 1.0 / ray.direction()[i]; + let mut t0 = (self.min[i] - ray.origin()[i]) * inverse_dir; + let mut t1 = (self.max[i] - ray.origin()[i]) * inverse_dir; + if inverse_dir < 0.0 { + std::mem::swap(&mut t0, &mut t1); + } + t_min = if t0 > t_min { t0 } else { t_min }; + t_max = if t1 < t_max { t1 } else { t_max }; + + if t_max <= t_min { + return false; + } + } + + true + } + + pub fn surrounding_box(box0: Aabb, box1: Aabb) -> Self { + let smol_box = Vec3::new( + box0.min.x().min(box1.min.x()), + box0.min.y().min(box1.min.y()), + box0.min.z().min(box1.min.z()), + ); + + let big_box = Vec3::new( + box0.max.x().max(box1.max.x()), + box0.max.y().max(box1.max.y()), + box0.max.z().max(box1.max.z()), + ); + + Self { + min: smol_box, + max: big_box, + } + } +} diff --git a/src/bvh.rs b/src/bvh.rs new file mode 100644 index 0000000..a7bd00b --- /dev/null +++ b/src/bvh.rs @@ -0,0 +1,134 @@ +use std::cmp::Ordering; + +use rand::{prelude::SliceRandom, Rng}; + +use crate::{types::Ray, Aabb, HitRecord, Hitable}; + +pub struct BvhNode { + bounding_box: Aabb, + left: HitNode, + right: HitNode, +} + +impl BvhNode { + pub fn new(rng: &mut R, objects: &mut [T], t0: f64, t1: f64) -> Self { + let comparator = [ + Self::box_x_compare, + Self::box_y_compare, + Self::box_z_compare, + ] + .choose(rng) + .unwrap(); + + objects.sort_by(comparator); + + let (left, right) = match objects.len() { + 1 => ( + HitNode::Direct(objects[0].clone()), + HitNode::Direct(objects[0].clone()), + ), + 2 => match comparator(&objects[0], &objects[1]) { + Ordering::Greater => ( + HitNode::Direct(objects[1].clone()), + HitNode::Direct(objects[0].clone()), + ), + _ => ( + HitNode::Direct(objects[0].clone()), + HitNode::Direct(objects[1].clone()), + ), + }, + + n => { + let (l, r) = objects.split_at_mut(n / 2); + ( + HitNode::Bvh(Box::new(BvhNode::new(rng, l, t0, t1))), + HitNode::Bvh(Box::new(BvhNode::new(rng, r, t0, t1))), + ) + } + }; + + let left_box = left + .bounding_box(t0, t1) + .expect("missing bounding box for left BVH Node"); + let right_box = right + .bounding_box(t0, t1) + .expect("missing bounding box for right BVH Node"); + + Self { + left, + right, + bounding_box: Aabb::surrounding_box(left_box, right_box), + } + } + + fn box_x_compare(obj1: &T, obj2: &T) -> Ordering { + if let (Some(bbox_a), Some(bbox_b)) = + (obj1.bounding_box(0.0, 0.0), obj2.bounding_box(0.0, 0.0)) + { + return bbox_a.min.x().partial_cmp(&bbox_b.min.x()).unwrap(); + } + + panic!("No bounding box for this BVH Node!!") + } + + fn box_y_compare(obj1: &T, obj2: &T) -> Ordering { + if let (Some(bbox_a), Some(bbox_b)) = + (obj1.bounding_box(0.0, 0.0), obj2.bounding_box(0.0, 0.0)) + { + return bbox_a.min.y().partial_cmp(&bbox_b.min.y()).unwrap(); + } + + panic!("No bounding box for this BVH Node!!") + } + + fn box_z_compare(obj1: &T, obj2: &T) -> Ordering { + if let (Some(bbox_a), Some(bbox_b)) = + (obj1.bounding_box(0.0, 0.0), obj2.bounding_box(0.0, 0.0)) + { + return bbox_a.min.z().partial_cmp(&bbox_b.min.z()).unwrap(); + } + + panic!("No bounding box for this BVH Node!!") + } +} + +impl Hitable for BvhNode { + fn hit(&self, ray: &Ray, t_min: f64, t_max: f64) -> Option { + self.bounding_box(t_min, t_max)?; + + let hbox_left = self.left.hit(ray, t_min, t_max); + + let hbox_right = if let Some(ref hleft) = hbox_left { + self.right.hit(ray, t_min, hleft.t) + } else { + self.right.hit(ray, t_min, t_max) + }; + + hbox_right.or(hbox_left) + } + + fn bounding_box(&self, _t_min: f64, _t_max: f64) -> Option { + Some(self.bounding_box) + } +} + +enum HitNode { + Bvh(Box>), + Direct(T), +} + +impl HitNode { + fn hit(&self, ray: &Ray, t_min: f64, t_max: f64) -> Option { + match self { + HitNode::Bvh(node) => node.hit(ray, t_min, t_max), + HitNode::Direct(node) => node.hit(ray, t_min, t_max), + } + } + + fn bounding_box(&self, t0: f64, t1: f64) -> Option { + match self { + HitNode::Bvh(node) => node.bounding_box(t0, t1), + HitNode::Direct(node) => node.bounding_box(t0, t1), + } + } +} diff --git a/src/demo.rs b/src/demo.rs index a3e745e..6154d62 100644 --- a/src/demo.rs +++ b/src/demo.rs @@ -1,10 +1,12 @@ +use crate::BvhNode; + use { crate::{ types::{ material::{Dielectric, Lambertian, Metal}, MovingSphere, Ray, Sphere, Vec3, }, - Camera, Hitable, HitableList, HORIZONTAL_PARTITION, VERTICAL_PARTITION, + Camera, Hitable, HORIZONTAL_PARTITION, VERTICAL_PARTITION, }, rand::{rngs::SmallRng, Rng, SeedableRng}, rayon::prelude::*, @@ -49,9 +51,7 @@ impl Demo { } fn world(&self) -> impl Hitable { - let mut world = HitableList { - list: Vec::with_capacity(500), - }; + let mut world: Vec> = Vec::with_capacity(500); let mut rng = rand::thread_rng(); let mut rng = SmallRng::from_rng(&mut rng).unwrap(); @@ -125,7 +125,7 @@ impl Demo { Metal::with_fuzz(Vec3::new(0.7, 0.6, 0.5), 0.0), ))); - world + BvhNode::new(&mut rng, &mut world, 0.0, 1.0) } fn camera(&self, aspect_ratio: f64) -> Camera { diff --git a/src/hitable.rs b/src/hitable.rs index 625fd24..3de389d 100644 --- a/src/hitable.rs +++ b/src/hitable.rs @@ -1,4 +1,9 @@ -use crate::types::{Material, Ray, Vec3}; +use std::sync::Arc; + +use crate::{ + types::{Material, Ray, Vec3}, + Aabb, +}; pub struct HitRecord<'a> { /// Rays are represented by A + t * B @@ -22,7 +27,16 @@ pub struct HitRecord<'a> { } pub trait Hitable { - fn hit(&self, _ray: &Ray, _t_min: f64, _t_max: f64) -> Option { - None + fn hit(&self, _ray: &Ray, _t_min: f64, _t_max: f64) -> Option; + + fn bounding_box(&self, _t0: f64, _t1: f64) -> Option; +} + +impl Hitable for Arc { + fn hit(&self, ray: &Ray, t_min: f64, t_max: f64) -> Option { + self.as_ref().hit(ray, t_min, t_max) + } + fn bounding_box(&self, t0: f64, t1: f64) -> Option { + self.as_ref().bounding_box(t0, t1) } } diff --git a/src/hitable_list.rs b/src/hitable_list.rs index d2f52cd..26fd415 100644 --- a/src/hitable_list.rs +++ b/src/hitable_list.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use crate::{demo::ParallelHit, types::Ray, HitRecord, Hitable}; +use crate::{demo::ParallelHit, types::Ray, Aabb, HitRecord, Hitable}; pub struct HitableList { pub list: Vec>, @@ -18,6 +18,28 @@ impl Hitable for HitableList { } hit_rec } + + fn bounding_box(&self, t0: f64, t1: f64) -> Option { + if self.list.is_empty() { + return None; + } + + let mut output_box = None; + + for obj in self.list.iter() { + if let Some(bbox) = obj.bounding_box(t0, t1) { + if let Some(ref mut opbox) = output_box { + *opbox = Aabb::surrounding_box(*opbox, bbox); + } else { + output_box = Some(bbox); + } + } else { + return output_box; + } + } + + output_box + } } impl HitableList { diff --git a/src/main.rs b/src/main.rs index 261fcc1..7b859b1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,7 @@ #![allow(clippy::suspicious_arithmetic_impl)] +mod aabb; +mod bvh; mod camera; mod demo; mod hitable; @@ -8,13 +10,15 @@ mod types; pub use camera::Camera; +pub use aabb::Aabb; +pub use bvh::BvhNode; pub use hitable::{HitRecord, Hitable}; pub use hitable_list::HitableList; use std::time::Instant; -const NUM_SAMPLES: u8 = 5; -const VERTICAL_PARTITION: usize = 8; -const HORIZONTAL_PARTITION: usize = 8; +const NUM_SAMPLES: u8 = 25; +const VERTICAL_PARTITION: usize = 12; +const HORIZONTAL_PARTITION: usize = 12; const WIDTH: usize = 1920; const HEIGHT: usize = 1080; diff --git a/src/types/moving_sphere.rs b/src/types/moving_sphere.rs index c680a9e..a866ae6 100644 --- a/src/types/moving_sphere.rs +++ b/src/types/moving_sphere.rs @@ -1,6 +1,6 @@ use crate::{ types::{Material, Ray, Vec3}, - HitRecord, Hitable, + Aabb, HitRecord, Hitable, }; pub struct MovingSphere { @@ -72,4 +72,12 @@ impl Hitable for MovingSphere { } None } + + fn bounding_box(&self, t0: f64, t1: f64) -> Option { + let radius = Vec3::new(self.radius, self.radius, self.radius); + let box_smol = Aabb::new(self.center(t0) - radius, self.center(t0) + radius); + let box_big = Aabb::new(self.center(t1) - radius, self.center(t1) + radius); + + Some(Aabb::surrounding_box(box_smol, box_big)) + } } diff --git a/src/types/sphere.rs b/src/types/sphere.rs index 1df8d74..a915eba 100644 --- a/src/types/sphere.rs +++ b/src/types/sphere.rs @@ -1,6 +1,6 @@ use crate::{ types::{Material, Ray, Vec3}, - HitRecord, Hitable, + Aabb, HitRecord, Hitable, }; pub struct Sphere { @@ -61,4 +61,9 @@ impl Hitable for Sphere { } None } + + fn bounding_box(&self, _t0: f64, _t1: f64) -> Option { + let radius = Vec3::new(self.radius, self.radius, self.radius); + Some(Aabb::new(self.center - radius, self.center + radius)) + } } diff --git a/src/types/vec3.rs b/src/types/vec3.rs index 12d931e..6235253 100644 --- a/src/types/vec3.rs +++ b/src/types/vec3.rs @@ -3,7 +3,7 @@ use std::{ ops::{Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign}, }; -#[derive(Copy, Clone)] +#[derive(Debug, Copy, Clone)] pub struct Vec3([f64; 3]); impl Vec3 {