rustracer/src/bvh.rs

use crate::hittable::{HittableList, Hittable, HitRecord, HittableObject};
use crate::aabb::Aabb;
use std::cmp;
use crate::Ray;

// https://github.com/fralken/ray-tracing-the-next-week/blob/master/src/bvh.rs

enum BVHNode {
    Branch { left: Box<BVH>, right: Box<BVH> },
    Leaf(HittableObject)
}

pub struct BVH {
    tree: BVHNode,
    bounding_box: Aabb
}

impl BVH {
    pub fn new(mut objects: HittableList, time0: f64, time1: f64) -> Self {
        fn box_compare(time0: f64, time1: f64, axis: usize) -> impl FnMut(&HittableObject, &HittableObject) -> cmp::Ordering {
            move |a, b| {
                let a_bbox = a.bounding_box(time0, time1);
                let b_bbox = b.bounding_box(time0, time1);
                if let (Some(a), Some(b)) = (a_bbox, b_bbox) {
                    let ac = a.min()[axis] + a.max()[axis];
                    let bc = b.min()[axis] + b.max()[axis];
                    ac.partial_cmp(&bc).unwrap()
                } else {
                    panic!("no bounding box in bvh node")
                }
            }
        }

        fn axis_range(objects: &HittableList, time0: f64, time1: f64, axis: usize) -> f64 {
            let (min, max) = objects.iter().fold((f64::MAX, f64::MIN), |(bmin, bmax), hit| {
                if let Some(aabb) = hit.bounding_box(time0, time1) {
                    (bmin.min(aabb.min()[axis]), bmax.max(aabb.max()[axis]))
                } else {
                    (bmin, bmax)
                }
            });
            max - min
        }

        let mut axis_ranges: Vec<(usize, f64)> = (0..3)
            .map(|a| (a, axis_range(&objects, time0, time1, a)))
            .collect();

        axis_ranges.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

        let axis = axis_ranges[0].0;

        objects.sort_unstable_by(box_compare(time0, time1, axis));
        let len = objects.len();
        match len {
            0 => panic!("no elements in scene"),
            1 => {
                let leaf = objects.pop().unwrap();
                if let Some(bounding_box) = leaf.bounding_box(time0, time1) {
                    BVH { tree: BVHNode::Leaf(leaf), bounding_box }
                } else {
                    panic!("no bounding box in bvh node")
                }
            },
            _ => {
                let right = BVH::new(objects.drain(len / 2..).collect(), time0, time1);
                let left = BVH::new(objects, time0, time1);
                let bounding_box = left.bounding_box.surrounding(&right.bounding_box);
                BVH { tree: BVHNode::Branch { left: Box::new(left), right: Box::new(right) }, bounding_box }
            }
        }
    }
}

impl Hittable for BVH {
    fn hit(&self, ray: &Ray, t_min: f64, mut t_max: f64) -> Option<HitRecord> {
        if !self.bounding_box.hit(ray, t_min, t_max) {
            return None;
        }
        match &self.tree {
            BVHNode::Leaf(leaf) => leaf.hit(ray, t_min, t_max),
            BVHNode::Branch { left, right } => {
                let hit_left = left.hit(ray, t_min, t_max);
                if let Some(hl) = &hit_left { t_max = hl.t };
                let hit_right = right.hit(ray, t_min, t_max);
                if hit_right.is_some() { hit_right } else { hit_left }
            }
        }
    }
    fn bounding_box(&self, _time0: f64, _time1: f64) -> Option<Aabb> {
        Some(self.bounding_box)
    }
}