use crate::bounding_volume::{Aabb, SimdAabb};
use crate::math::Vector;
use crate::math::{Point, Real};
use crate::query::SplitResult;
use crate::simd::SimdReal;
use simba::simd::SimdValue;
use super::utils::split_indices_wrt_dim;
use super::{IndexedData, NodeIndex, Qbvh, QbvhNode, QbvhNodeFlags, QbvhProxy};
pub struct BuilderProxies<'a, LeafData> {
proxies: &'a mut Vec<QbvhProxy<LeafData>>,
aabbs: &'a mut Vec<Aabb>,
}
impl<'a, LeafData> BuilderProxies<'a, LeafData> {
fn insert(&mut self, data: LeafData, aabb: Aabb)
where
LeafData: IndexedData,
{
let index = data.index();
if self.proxies.len() <= index {
self.proxies.resize(index + 1, QbvhProxy::invalid());
self.aabbs.resize(index + 1, Aabb::new_invalid());
}
self.proxies[index] = QbvhProxy::detached(data);
self.aabbs[index] = aabb;
}
}
pub trait QbvhDataSplitter<LeafData> {
fn split_dataset<'idx>(
&mut self,
subdiv_dims: [usize; 2],
center: Point<Real>,
indices: &'idx mut [usize],
indices_workspace: &'idx mut Vec<usize>,
proxies: BuilderProxies<LeafData>,
) -> [&'idx mut [usize]; 4];
}
pub struct CenterDataSplitter {
pub enable_fallback_split: bool,
}
impl Default for CenterDataSplitter {
fn default() -> Self {
Self {
enable_fallback_split: true,
}
}
}
impl<LeafData> QbvhDataSplitter<LeafData> for CenterDataSplitter {
fn split_dataset<'idx>(
&mut self,
subdiv_dims: [usize; 2],
center: Point<Real>,
indices: &'idx mut [usize],
_: &'idx mut Vec<usize>,
proxies: BuilderProxies<LeafData>,
) -> [&'idx mut [usize]; 4] {
self.split_dataset_wo_workspace(subdiv_dims, center, indices, &*proxies.aabbs)
}
}
impl CenterDataSplitter {
pub(crate) fn split_dataset_wo_workspace<'idx>(
&self,
subdiv_dims: [usize; 2],
center: Point<Real>,
indices: &'idx mut [usize],
aabbs: &[Aabb],
) -> [&'idx mut [usize]; 4] {
let (left, right) = split_indices_wrt_dim(
indices,
aabbs,
¢er,
subdiv_dims[0],
self.enable_fallback_split,
);
let (left_bottom, left_top) = split_indices_wrt_dim(
left,
aabbs,
¢er,
subdiv_dims[1],
self.enable_fallback_split,
);
let (right_bottom, right_top) = split_indices_wrt_dim(
right,
aabbs,
¢er,
subdiv_dims[1],
self.enable_fallback_split,
);
[left_bottom, left_top, right_bottom, right_top]
}
}
pub struct QbvhNonOverlappingDataSplitter<F> {
pub canonical_split: F,
pub epsilon: Real,
}
impl<LeafData, F> QbvhDataSplitter<LeafData> for QbvhNonOverlappingDataSplitter<F>
where
LeafData: IndexedData,
F: FnMut(LeafData, usize, Real, Real, Aabb, Aabb) -> SplitResult<(LeafData, Aabb)>,
{
fn split_dataset<'idx>(
&mut self,
subdiv_dims: [usize; 2],
center: Point<Real>,
indices: &'idx mut [usize],
indices_workspace: &'idx mut Vec<usize>,
mut proxies: BuilderProxies<LeafData>,
) -> [&'idx mut [usize]; 4] {
let mut split_pt = Point::from(Vector::repeat(-Real::MAX));
let mut split_pt_right = Point::from(Vector::repeat(Real::MAX));
for dim in subdiv_dims {
for i in indices.iter().copied() {
let aabb = &proxies.aabbs[i];
if aabb.maxs[dim] <= center[dim] && aabb.maxs[dim] > split_pt[dim] {
split_pt[dim] = aabb.maxs[dim];
}
if aabb.mins[dim] >= center[dim] && aabb.mins[dim] < split_pt_right[dim] {
split_pt_right[dim] = aabb.mins[dim];
}
}
if (split_pt[dim] - center[dim]).abs() > (split_pt_right[dim] - center[dim]).abs() {
split_pt[dim] = split_pt_right[dim];
}
if split_pt[dim] == -Real::MAX || split_pt[dim] == Real::MAX {
let candidate_min = proxies.aabbs[indices[0]].mins[dim];
let candidate_max = proxies.aabbs[indices[0]].maxs[dim];
for i in indices.iter().copied() {
let aabb = &proxies.aabbs[i];
if aabb.mins[dim] < candidate_min {
split_pt[dim] = candidate_min;
break;
} else if aabb.mins[dim] > candidate_min {
split_pt[dim] = aabb.mins[dim];
}
if aabb.maxs[dim] > candidate_max {
split_pt[dim] = candidate_max;
break;
} else if aabb.maxs[dim] < candidate_max {
split_pt[dim] = aabb.maxs[dim];
}
}
}
}
if (split_pt[subdiv_dims[0]] == -Real::MAX || split_pt[subdiv_dims[0]] == Real::MAX)
&& (split_pt[subdiv_dims[1]] == -Real::MAX || split_pt[subdiv_dims[1]] == Real::MAX)
{
split_pt = center;
}
indices_workspace.resize(indices.len(), 0);
indices_workspace.copy_from_slice(indices);
for dim in subdiv_dims {
for k in 0..indices_workspace.len() {
let i = indices_workspace[k];
if let SplitResult::Pair(aabb_l, aabb_r) =
proxies.aabbs[i].canonical_split(dim, split_pt[dim], self.epsilon)
{
if let SplitResult::Pair((data_l, aabb_l), (data_r, aabb_r)) = (self
.canonical_split)(
proxies.proxies[i].data,
dim,
split_pt[dim],
self.epsilon,
aabb_l,
aabb_r,
) {
indices_workspace[k] = data_l.index();
indices_workspace.push(data_r.index());
proxies.insert(data_l, aabb_l);
proxies.insert(data_r, aabb_r);
}
}
}
}
let center_splitter = CenterDataSplitter {
enable_fallback_split: false,
};
center_splitter.split_dataset_wo_workspace(
subdiv_dims,
split_pt,
indices_workspace,
&*proxies.aabbs,
)
}
}
pub trait QbvhDataGenerator<LeafData> {
fn size_hint(&self) -> usize;
fn for_each(&mut self, f: impl FnMut(LeafData, Aabb));
}
impl<LeafData, F> QbvhDataGenerator<LeafData> for F
where
F: ExactSizeIterator<Item = (LeafData, Aabb)>,
{
fn size_hint(&self) -> usize {
self.len()
}
#[inline(always)]
fn for_each(&mut self, mut f: impl FnMut(LeafData, Aabb)) {
for (elt, aabb) in self {
f(elt, aabb)
}
}
}
impl<LeafData: IndexedData> Qbvh<LeafData> {
pub fn clear_and_rebuild(
&mut self,
data_gen: impl QbvhDataGenerator<LeafData>,
dilation_factor: Real,
) {
self.clear_and_rebuild_with_splitter(
data_gen,
CenterDataSplitter::default(),
dilation_factor,
);
}
}
impl<LeafData: IndexedData> Qbvh<LeafData> {
pub fn clear_and_rebuild_with_splitter(
&mut self,
mut data_gen: impl QbvhDataGenerator<LeafData>,
mut splitter: impl QbvhDataSplitter<LeafData>,
dilation_factor: Real,
) {
self.free_list.clear();
self.nodes.clear();
self.proxies.clear();
let mut indices = Vec::with_capacity(data_gen.size_hint());
let mut aabbs = vec![Aabb::new_invalid(); data_gen.size_hint()];
self.proxies = vec![QbvhProxy::invalid(); data_gen.size_hint()];
data_gen.for_each(|data, aabb| {
let index = data.index();
if index >= self.proxies.len() {
self.proxies.resize(index + 1, QbvhProxy::invalid());
aabbs.resize(index + 1, Aabb::new_invalid());
}
self.proxies[index].data = data;
aabbs[index] = aabb;
indices.push(index);
});
let root_node = QbvhNode {
simd_aabb: SimdAabb::new_invalid(),
children: [1, u32::MAX, u32::MAX, u32::MAX],
parent: NodeIndex::invalid(),
flags: QbvhNodeFlags::default(),
};
self.nodes.push(root_node);
let root_id = NodeIndex::new(0, 0);
let (_, aabb) = self.do_recurse_build_generic(
&mut splitter,
&mut indices,
&mut aabbs,
root_id,
dilation_factor,
);
self.root_aabb = aabb;
self.nodes[0].simd_aabb = SimdAabb::from([
aabb,
Aabb::new_invalid(),
Aabb::new_invalid(),
Aabb::new_invalid(),
]);
}
fn do_recurse_build_generic(
&mut self,
splitter: &mut impl QbvhDataSplitter<LeafData>,
indices: &mut [usize],
aabbs: &mut Vec<Aabb>,
parent: NodeIndex,
dilation: Real,
) -> (u32, Aabb) {
if indices.len() <= 4 {
let my_id = self.nodes.len();
let mut leaf_aabbs = [Aabb::new_invalid(); 4];
let mut proxy_ids = [u32::MAX; 4];
for (k, id) in indices.iter().enumerate() {
leaf_aabbs[k] = aabbs[*id];
proxy_ids[k] = *id as u32;
self.proxies[*id].node = NodeIndex::new(my_id as u32, k as u8);
}
let mut node = QbvhNode {
simd_aabb: SimdAabb::from(leaf_aabbs),
children: proxy_ids,
parent,
flags: QbvhNodeFlags::LEAF,
};
node.simd_aabb.dilate_by_factor(SimdReal::splat(dilation));
let my_aabb = node.simd_aabb.to_merged_aabb();
self.nodes.push(node);
return (my_id as u32, my_aabb);
}
let mut center = Point::origin();
#[cfg(feature = "dim3")]
let mut variance = Vector::zeros();
let center_denom = 1.0 / (indices.len() as Real);
for i in &*indices {
let coords = aabbs[*i].center().coords;
center += coords * center_denom;
}
#[cfg(feature = "dim3")]
{
let variance_denom = 1.0 / ((indices.len() - 1) as Real);
for i in &*indices {
let dir_to_center = aabbs[*i].center() - center;
variance += dir_to_center.component_mul(&dir_to_center) * variance_denom;
}
}
#[allow(unused_mut)] let mut subdiv_dims = [0, 1];
#[cfg(feature = "dim3")]
{
let min = variance.imin();
subdiv_dims[0] = (min + 1) % 3;
subdiv_dims[1] = (min + 2) % 3;
}
let node = QbvhNode {
simd_aabb: SimdAabb::new_invalid(),
children: [0; 4], parent,
flags: QbvhNodeFlags::default(),
};
let id = self.nodes.len() as u32;
self.nodes.push(node);
let proxies = BuilderProxies {
proxies: &mut self.proxies,
aabbs,
};
let mut workspace = vec![]; let splits = splitter.split_dataset(subdiv_dims, center, indices, &mut workspace, proxies);
let n = [
NodeIndex::new(id, 0),
NodeIndex::new(id, 1),
NodeIndex::new(id, 2),
NodeIndex::new(id, 3),
];
let children = [
self.do_recurse_build_generic(splitter, splits[0], aabbs, n[0], dilation),
self.do_recurse_build_generic(splitter, splits[1], aabbs, n[1], dilation),
self.do_recurse_build_generic(splitter, splits[2], aabbs, n[2], dilation),
self.do_recurse_build_generic(splitter, splits[3], aabbs, n[3], dilation),
];
self.nodes[id as usize].children =
[children[0].0, children[1].0, children[2].0, children[3].0];
self.nodes[id as usize].simd_aabb =
SimdAabb::from([children[0].1, children[1].1, children[2].1, children[3].1]);
self.nodes[id as usize]
.simd_aabb
.dilate_by_factor(SimdReal::splat(dilation));
let my_aabb = self.nodes[id as usize].simd_aabb.to_merged_aabb();
(id, my_aabb)
}
}