raytracers/rtiow/renderer/src/kdtree.rs

use std::fmt;

use log::info;

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

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

// 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!(
            "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 KDTree {
    pub fn new(l: Vec<Box<dyn Hit>>, t_min: f32, t_max: f32) -> KDTree {
        if l.is_empty() {
            panic!("Attempt to build k-d tree with no Hit objects");
        }

        if l.len() == 1 {
            return KDTree::Leaf(vec_first_into(l));
        }

        if l.len() == 2 {
            let (left_vec, right_vec) = vec_split_into(l, 1);
            let left = vec_first_into(left_vec);
            let right = vec_first_into(right_vec);
            let bbox = surrounding_box(
                &left.bounding_box(t_min, t_max).unwrap(),
                &right.bounding_box(t_min, t_max).unwrap(),
            );
            return KDTree::Branch {
                left: Box::new(KDTree::Leaf(left)),
                right: Box::new(KDTree::Leaf(right)),
                bbox: Some(bbox),
            };
        }
        let bbox = l[0].bounding_box(t_min, t_max).unwrap();
        let mut mid_point = bbox.mid_point();
        let bbox = l[1..].iter().fold(bbox, |acc, h| {
            let new_bbox = &h.bounding_box(t_min, t_max).unwrap();
            mid_point = mid_point + new_bbox.mid_point();
            surrounding_box(&acc, new_bbox)
        });
        mid_point = mid_point / l.len() as f32;
        let axis = bbox.longest_axis();
        let mut left_half = Vec::new();
        let mut right_half = Vec::new();
        match axis {
            0 => l.into_iter().for_each(|h| {
                if h.bounding_box(t_min, t_max).unwrap().mid_point().x < mid_point.x {
                    left_half.push(h);
                } else {
                    right_half.push(h);
                }
            }),
            1 => l.into_iter().for_each(|h| {
                if h.bounding_box(t_min, t_max).unwrap().mid_point().y < mid_point.y {
                    left_half.push(h);
                } else {
                    right_half.push(h);
                }
            }),
            2 => l.into_iter().for_each(|h| {
                if h.bounding_box(t_min, t_max).unwrap().mid_point().z < mid_point.z {
                    left_half.push(h);
                } else {
                    right_half.push(h);
                }
            }),
            _ => panic!("Unreachable"),
        };
        KDTree::Branch {
            left: Box::new(KDTree::new(left_half, t_min, t_max)),
            right: Box::new(KDTree::new(right_half, t_min, t_max)),
            bbox: Some(bbox),
        }
    }

    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()
    }
}

fn print_tree(f: &mut fmt::Formatter, depth: usize, kdt: &KDTree) -> fmt::Result {
    let vol = kdt.volume();
    write!(f, "{:8.2}{}", vol, " ".repeat(depth * 2))?;
    match kdt {
        KDTree::Leaf(ref hit) => writeln!(
            f,
            "Leaf: {}",
            hit.bounding_box(0., 0.)
                .map_or("NO BBOX".to_owned(), |bb| bb.to_string())
        )?,
        KDTree::Branch { bbox, .. } => writeln!(
            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())
        )?,
    }
    if let KDTree::Branch { left, right, .. } = kdt {
        print_tree(f, depth + 1, left)?;
        print_tree(f, depth + 1, right)?;
    }
    Ok(())
}

impl fmt::Display for KDTree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "kd-tree")?;
        print_tree(f, 1, self)
    }
}
impl Hit for KDTree {
    fn hit(&self, r: Ray, t_min: f32, t_max: f32) -> Option<HitRecord> {
        match self.bounding_box(t_min, t_max) {
            Some(bbox) => {
                if !bbox.hit(r, t_min, t_max) {
                    return None;
                }
            }
            None => {
                info!("KDTree::hit no bbox: {:?}", r);
                return None;
            }
        }
        match self {
            KDTree::Leaf(ref hit) => hit.hit(r, t_min, t_max),
            KDTree::Branch { left, right, bbox } => match bbox {
                None => None,
                Some(_bbox) => 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,
                },
            },
        }
    }

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