raytracers/rtiow/src/bvh.rs

use std;
use std::fmt;
use std::time::Instant;

use rand;
use rand::Rng;

use crate::aabb::surrounding_box;
use crate::aabb::AABB;
use crate::hitable::Hit;
use crate::hitable::HitRecord;
use crate::ray::Ray;

enum BVHNode {
    Leaf(Box<Hit>),
    Branch {
        left: Box<BVHNode>,
        right: Box<BVHNode>,
        bbox: Option<AABB>,
    },
}

impl fmt::Display for BVHNode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            BVHNode::Leaf(ref hit) => write!(
                f,
                "Leaf: {}",
                hit.bounding_box(0., 0.)
                    .map_or("NO BBOX".to_owned(), |bb| bb.to_string())
            ),
            BVHNode::Branch { bbox, .. } => write!(
                f,
                "Branch: {}",
                // TODO(wathiede): removing this .clone() results in a complaint about moving out
                // of a borrow.
                bbox.clone()
                    .map_or("NO BBOX".to_owned(), |bb| bb.to_string())
            ),
        }
    }
}

// Lint is wrong when dealing with Box<Trait>
#[cfg_attr(feature = "cargo-clippy", allow(borrowed_box))]
fn box_x_compare(ah: &Box<Hit>, bh: &Box<Hit>) -> std::cmp::Ordering {
    match (ah.bounding_box(0., 0.), bh.bounding_box(0., 0.)) {
        (Some(box_left), Some(box_right)) => {
            box_left.min().x.partial_cmp(&box_right.min().x).unwrap()
        }
        _ => panic!("hit missing bounding box"),
    }
}

#[cfg_attr(feature = "cargo-clippy", allow(borrowed_box))]
fn box_y_compare(ah: &Box<Hit>, bh: &Box<Hit>) -> std::cmp::Ordering {
    match (ah.bounding_box(0., 0.), bh.bounding_box(0., 0.)) {
        (Some(box_left), Some(box_right)) => {
            box_left.min().y.partial_cmp(&box_right.min().y).unwrap()
        }
        _ => panic!("hit missing bounding box"),
    }
}

#[cfg_attr(feature = "cargo-clippy", allow(borrowed_box))]
fn box_z_compare(ah: &Box<Hit>, bh: &Box<Hit>) -> std::cmp::Ordering {
    match (ah.bounding_box(0., 0.), bh.bounding_box(0., 0.)) {
        (Some(box_left), Some(box_right)) => {
            box_left.min().z.partial_cmp(&box_right.min().z).unwrap()
        }
        _ => panic!("hit missing bounding box"),
    }
}

// Return the first element from the vector, which should be the only element.
// TODO(wathiede): we really want a .first_into() (.first() doesn't work because it
// returns a reference.)
fn vec_first_into<T>(v: Vec<T>) -> T {
    if v.len() != 1 {
        panic!(format!(
            "vec_first_into called for vector length != 1, length {}",
            v.len()
        ));
    }
    v.into_iter().next().unwrap()
}

fn vec_split_into<T>(v: Vec<T>, offset: usize) -> (Vec<T>, Vec<T>) {
    let mut left_half = Vec::new();
    let mut right_half = Vec::new();
    v.into_iter().enumerate().for_each(|(i, h)| {
        if i < offset {
            left_half.push(h);
        } else {
            right_half.push(h);
        }
    });
    (left_half, right_half)
}

impl BVHNode {
    fn new(l: Vec<Box<Hit>>, t_min: f32, t_max: f32) -> BVHNode {
        if l.len() == 1 {
            BVHNode::Leaf(vec_first_into(l))
        } else {
            let mut l: Vec<Box<Hit>> = l.into_iter().collect();
            let mut rng = rand::thread_rng();
            match rng.gen_range::<u16>(0, 3) {
                0 => l.sort_by(box_x_compare),
                1 => l.sort_by(box_y_compare),
                2 => l.sort_by(box_z_compare),
                val => panic!("unknown axis {}", val),
            }

            let half = l.len() / 2;
            let (left_half, right_half) = vec_split_into(l, half);
            let left = Box::new(BVHNode::new(left_half, t_min, t_max));
            let right = Box::new(BVHNode::new(right_half, t_min, t_max));
            let bbox = BVHNode::surrounding_box(left.as_ref(), right.as_ref(), t_min, t_max);
            BVHNode::Branch { left, right, bbox }
        }
    }

    fn surrounding_box(left: &Hit, right: &Hit, t_min: f32, t_max: f32) -> Option<AABB> {
        match (
            left.bounding_box(t_min, t_max),
            right.bounding_box(t_min, t_max),
        ) {
            (Some(left_bbox), Some(right_bbox)) => Some(surrounding_box(&left_bbox, &right_bbox)),
            (None, Some(right_bbox)) => Some(right_bbox),
            (Some(left_bbox), None) => Some(left_bbox),
            (None, None) => None,
        }
    }

    fn volume(&self) -> f32 {
        let bbox = self.bounding_box(0., 0.).unwrap();
        (bbox.min().x - bbox.max().x).abs()
            * (bbox.min().y - bbox.max().y).abs()
            * (bbox.min().z - bbox.max().z).abs()
    }
}

impl Hit for BVHNode {
    fn hit(&self, r: Ray, t_min: f32, t_max: f32) -> Option<HitRecord> {
        match self {
            BVHNode::Leaf(ref hit) => hit.hit(r, t_min, t_max),
            BVHNode::Branch {
                ref left,
                ref right,
                ref bbox,
            } => match bbox {
                Some(ref bbox) => {
                    if bbox.hit(r, t_min, t_max) {
                        match (left.hit(r, t_min, t_max), right.hit(r, t_min, t_max)) {
                            (Some(hit_left), Some(hit_right)) => if hit_left.t < hit_right.t {
                                Some(hit_left)
                            } else {
                                Some(hit_right)
                            },
                            (Some(hit_left), None) => Some(hit_left),
                            (None, Some(hit_right)) => Some(hit_right),
                            (None, None) => None,
                        };
                    }
                    None
                }
                None => None,
            },
        }
    }

    fn bounding_box(&self, t_min: f32, t_max: f32) -> Option<AABB> {
        match self {
            BVHNode::Leaf(ref hit) => hit.bounding_box(t_min, t_max),
            BVHNode::Branch { ref bbox, .. } => bbox.clone(),
        }
    }
}

pub struct BVH {
    root: BVHNode,
}

impl BVH {
    pub fn new(l: Vec<Box<Hit>>, t_min: f32, t_max: f32) -> BVH {
        let count = l.len();
        let start = Instant::now();
        let bvh = BVH {
            root: BVHNode::new(l, t_min, t_max),
        };
        let runtime = start.elapsed();
        info!(
            "BVH build time {}.{} seconds for {} hitables",
            runtime.as_secs(),
            runtime.subsec_millis(),
            count
        );
        bvh
    }
}

fn print_tree(f: &mut fmt::Formatter, depth: usize, bvhn: &BVHNode) -> fmt::Result {
    let vol = bvhn.volume();
    writeln!(f, "{:.*}{}{}", 2, vol, " ".repeat(depth * 2), bvhn)?;
    if let BVHNode::Branch { left, right, .. } = bvhn {
        print_tree(f, depth + 1, left)?;
        print_tree(f, depth + 1, right)?;
    }
    Ok(())
}

impl fmt::Display for BVH {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "Root")?;
        print_tree(f, 1, &self.root)
    }
}

impl Hit for BVH {
    fn hit(&self, r: Ray, t_min: f32, t_max: f32) -> Option<HitRecord> {
        self.root.hit(r, t_min, t_max)
    }

    fn bounding_box(&self, t_min: f32, t_max: f32) -> Option<AABB> {
        self.root.bounding_box(t_min, t_max)
    }
}

#[cfg(test)]
mod tests {
    use crate::aabb::AABB;
    use crate::material::Lambertian;
    use crate::material::Metal;
    use crate::sphere::Sphere;
    use crate::texture::ConstantTexture;
    use crate::vec3::Vec3;

    use super::*;

    #[test]
    fn bbox_two_spheres() {
        let two_spheres_bvh = BVH::new(
            vec![
                Box::new(Sphere::new(
                    Vec3::new(0., 0., 0.),
                    0.5,
                    Lambertian::new(ConstantTexture::new(Vec3::new(0.1, 0.2, 0.5))),
                )),
                Box::new(Sphere::new(
                    Vec3::new(1., 0., 0.),
                    0.5,
                    Metal::new(Vec3::new(0.6, 0.6, 0.6), 0.2),
                )),
            ],
            0.,
            1.,
        );
        assert_eq!(
            AABB::new(Vec3::new(-0.5, -0.5, -0.5), Vec3::new(1.5, 0.5, 0.5)),
            two_spheres_bvh.bounding_box(0., 1.).unwrap(),
            "BVH:\n{}",
            two_spheres_bvh
        );
    }

    #[test]
    fn bbox_three_spheres() {
        let three_spheres_bvh = BVH::new(
            vec![
                Box::new(Sphere::new(
                    Vec3::new(0., 0., 0.),
                    0.5,
                    Lambertian::new(ConstantTexture::new(Vec3::new(0.1, 0.2, 0.5))),
                )),
                Box::new(Sphere::new(
                    Vec3::new(1., 0., 0.),
                    0.5,
                    Metal::new(Vec3::new(0.6, 0.6, 0.6), 0.2),
                )),
                Box::new(Sphere::new(
                    Vec3::new(0., 1., 0.),
                    0.5,
                    Metal::new(Vec3::new(0.6, 0.6, 0.6), 0.2),
                )),
            ],
            0.,
            1.,
        );
        assert_eq!(
            AABB::new(Vec3::new(-0.5, -0.5, -0.5), Vec3::new(1.5, 1.5, 0.5)),
            three_spheres_bvh.bounding_box(0., 1.).unwrap(),
            "BVH:\n{}",
            three_spheres_bvh
        );
    }
}