use std::iter::FusedIterator;
use std::mem;
use approx::abs_diff_eq;
use rayon::iter::{
IndexedParallelIterator, IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator,
};
use vek::Aabb;
#[derive(Clone)]
pub struct Bvh<T> {
internal_nodes: Vec<InternalNode>,
leaf_nodes: Vec<LeafNode<T>>,
root: NodeIdx,
}
#[derive(Clone)]
struct InternalNode {
bb: Aabb<f64>,
left: NodeIdx,
right: NodeIdx,
}
#[derive(Clone)]
struct LeafNode<T> {
bb: Aabb<f64>,
data: T,
}
type NodeIdx = u32;
impl<T: Send + Sync> Bvh<T> {
pub fn new() -> Self {
Self {
internal_nodes: Vec::new(),
leaf_nodes: Vec::new(),
root: NodeIdx::MAX,
}
}
pub fn build(&mut self, leaves: impl IntoIterator<Item = (T, Aabb<f64>)>) {
self.leaf_nodes.clear();
self.internal_nodes.clear();
self.leaf_nodes
.extend(leaves.into_iter().map(|(id, bb)| LeafNode { bb, data: id }));
let leaf_count = self.leaf_nodes.len();
if leaf_count == 0 {
return;
}
self.internal_nodes.reserve_exact(leaf_count - 1);
self.internal_nodes.resize(
leaf_count - 1,
InternalNode {
bb: Aabb::default(),
left: NodeIdx::MAX,
right: NodeIdx::MAX,
},
);
if NodeIdx::try_from(leaf_count)
.ok()
.and_then(|count| count.checked_add(count - 1))
.is_none()
{
panic!("too many elements in BVH");
}
let id = self.leaf_nodes[0].bb;
let scene_bounds = self
.leaf_nodes
.par_iter()
.map(|l| l.bb)
.reduce(|| id, Aabb::union);
self.root = build_rec(
0,
scene_bounds,
&mut self.internal_nodes,
&mut self.leaf_nodes,
leaf_count as NodeIdx,
)
.0;
debug_assert_eq!(self.internal_nodes.len(), self.leaf_nodes.len() - 1);
}
pub fn traverse(&self) -> Option<Node<T>> {
if !self.leaf_nodes.is_empty() {
Some(Node::from_idx(self, self.root))
} else {
None
}
}
pub fn iter(
&self,
) -> impl ExactSizeIterator<Item = (&T, Aabb<f64>)> + FusedIterator + Clone + '_ {
self.leaf_nodes.iter().map(|leaf| (&leaf.data, leaf.bb))
}
#[allow(dead_code)]
pub fn iter_mut(
&mut self,
) -> impl ExactSizeIterator<Item = (&mut T, Aabb<f64>)> + FusedIterator + '_ {
self.leaf_nodes
.iter_mut()
.map(|leaf| (&mut leaf.data, leaf.bb))
}
pub fn par_iter(&self) -> impl IndexedParallelIterator<Item = (&T, Aabb<f64>)> + Clone + '_ {
self.leaf_nodes.par_iter().map(|leaf| (&leaf.data, leaf.bb))
}
#[allow(dead_code)]
pub fn par_iter_mut(
&mut self,
) -> impl IndexedParallelIterator<Item = (&mut T, Aabb<f64>)> + '_ {
self.leaf_nodes
.par_iter_mut()
.map(|leaf| (&mut leaf.data, leaf.bb))
}
}
pub enum Node<'a, T> {
Internal(Internal<'a, T>),
Leaf { data: &'a T, bb: Aabb<f64> },
}
impl<'a, T> Node<'a, T> {
fn from_idx(bvh: &'a Bvh<T>, idx: NodeIdx) -> Self {
if idx < bvh.internal_nodes.len() as NodeIdx {
Self::Internal(Internal { bvh, idx })
} else {
let leaf = &bvh.leaf_nodes[(idx - bvh.internal_nodes.len() as NodeIdx) as usize];
Self::Leaf {
data: &leaf.data,
bb: leaf.bb,
}
}
}
pub fn bb(&self) -> Aabb<f64> {
match self {
Node::Internal(int) => int.bb(),
Node::Leaf { bb, .. } => *bb,
}
}
}
pub struct Internal<'a, T> {
bvh: &'a Bvh<T>,
idx: NodeIdx,
}
impl<'a, T> Internal<'a, T> {
pub fn split(self) -> (Aabb<f64>, Node<'a, T>, Node<'a, T>) {
let internal = &self.bvh.internal_nodes[self.idx as usize];
let bb = internal.bb;
let left = Node::from_idx(self.bvh, internal.left);
let right = Node::from_idx(self.bvh, internal.right);
(bb, left, right)
}
pub fn bb(&self) -> Aabb<f64> {
self.bvh.internal_nodes[self.idx as usize].bb
}
}
fn build_rec<T: Send>(
idx: NodeIdx,
mut bounds: Aabb<f64>,
internal_nodes: &mut [InternalNode],
leaf_nodes: &mut [LeafNode<T>],
total_leaf_count: NodeIdx,
) -> (NodeIdx, Aabb<f64>) {
debug_assert_eq!(leaf_nodes.len() - 1, internal_nodes.len());
if leaf_nodes.len() == 1 {
return (total_leaf_count - 1 + idx, leaf_nodes[0].bb);
}
loop {
debug_assert!(bounds.is_valid());
let dims = bounds.max - bounds.min;
let (mut split, bounds_left, bounds_right) = if dims.x >= dims.y && dims.x >= dims.z {
let mid = middle(bounds.min.x, bounds.max.x);
let [bounds_left, bounds_right] = bounds.split_at_x(mid);
let p = partition(leaf_nodes, |l| middle(l.bb.min.x, l.bb.max.x) <= mid);
(p, bounds_left, bounds_right)
} else if dims.y >= dims.x && dims.y >= dims.z {
let mid = middle(bounds.min.y, bounds.max.y);
let [bounds_left, bounds_right] = bounds.split_at_y(mid);
let p = partition(leaf_nodes, |l| middle(l.bb.min.y, l.bb.max.y) <= mid);
(p, bounds_left, bounds_right)
} else {
let mid = middle(bounds.min.z, bounds.max.z);
let [bounds_left, bounds_right] = bounds.split_at_z(mid);
let p = partition(leaf_nodes, |l| middle(l.bb.min.z, l.bb.max.z) <= mid);
(p, bounds_left, bounds_right)
};
if split == 0 {
if abs_diff_eq!(
bounds_right.min,
bounds_right.max,
epsilon = f64::EPSILON * 100.0
) {
split += 1;
} else {
bounds = bounds_right;
continue;
}
} else if split == leaf_nodes.len() {
if abs_diff_eq!(
bounds_left.min,
bounds_left.max,
epsilon = f64::EPSILON * 100.0
) {
split -= 1;
} else {
bounds = bounds_left;
continue;
}
}
let (leaves_left, leaves_right) = leaf_nodes.split_at_mut(split);
let (internal_left, internal_right) = internal_nodes.split_at_mut(split);
let (internal, internal_left) = internal_left.split_last_mut().unwrap();
let ((left, bounds_left), (right, bounds_right)) = rayon::join(
|| {
build_rec(
idx,
bounds_left,
internal_left,
leaves_left,
total_leaf_count,
)
},
|| {
build_rec(
idx + split as NodeIdx,
bounds_right,
internal_right,
leaves_right,
total_leaf_count,
)
},
);
internal.bb = bounds_left.union(bounds_right);
internal.left = left;
internal.right = right;
break (idx + split as NodeIdx - 1, internal.bb);
}
}
fn partition<T>(s: &mut [T], mut pred: impl FnMut(&T) -> bool) -> usize {
let mut it = s.iter_mut();
let mut true_count = 0;
while let Some(head) = it.find(|x| {
if pred(x) {
true_count += 1;
false
} else {
true
}
}) {
if let Some(tail) = it.rfind(|x| pred(x)) {
mem::swap(head, tail);
true_count += 1;
} else {
break;
}
}
true_count
}
fn middle(a: f64, b: f64) -> f64 {
(a + b) / 2.0
}
impl<T: Send + Sync> Default for Bvh<T> {
fn default() -> Self {
Self::new()
}
}