Implement AABB and BVH.

rtiow/src/bvh.rs

@@ -0,0 +1,149 @@
+use std;
+
+use rand;
+use rand::Rng;
+
+use aabb::surrounding_box;
+use aabb::AABB;
+use hitable::Hit;
+use hitable::HitRecord;
+use ray::Ray;
+
+enum BVHNode {
+    Leaf(Box<Hit>),
+    Branch {
+        left: Box<BVHNode>,
+        right: Box<BVHNode>,
+        bbox: Option<AABB>,
+    },
+}
+
+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)) => {
+            eprintln!("box_x_compare {} < {}", box_left.min, box_right.min);
+            return box_left.min.x.partial_cmp(&box_right.min.x).unwrap();
+        }
+        _ => panic!("hit missing bounding 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)) => {
+            eprintln!("box_y_compare {} < {}", box_left.min, box_right.min);
+            return box_left.min.y.partial_cmp(&box_right.min.y).unwrap();
+        }
+        _ => panic!("hit missing bounding 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)) => {
+            eprintln!("box_z_compare {} < {}", box_left.min, box_right.min);
+            return box_left.min.z.partial_cmp(&box_right.min.z).unwrap();
+        }
+        _ => panic!("hit missing bounding box"),
+    }
+}
+
+impl BVHNode {
+    fn new(l: Vec<Box<Hit>>, t_min: f32, t_max: f32) -> BVHNode {
+        if l.len() == 1 {
+            // 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.)
+            for h in l.into_iter() {
+                return BVHNode::Leaf(h);
+            }
+            panic!("Unreachable");
+        } 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 mut left_half = Vec::new();
+            let mut right_half = Vec::new();
+            let half_idx = l.len() / 2;
+            l.into_iter().enumerate().for_each(|(i, h)| {
+                if i < half_idx {
+                    left_half.push(h);
+                } else {
+                    right_half.push(h);
+                }
+            });
+            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);
+            return 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,
+        }
+    }
+}
+
+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,
+                ..
+            } => 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 {
+                    return Some(hit_left);
+                } else {
+                    return 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 {
+            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 {
+        BVH {
+            root: BVHNode::new(l, t_min, t_max),
+        }
+    }
+}
+
+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)
+    }
+}