use crate::bounding_volume::{Aabb, BoundingVolume, SimdAabb};
#[cfg(feature = "dim3")]
use crate::math::Vector;
use crate::math::{Point, Real};
use crate::partitioning::{CenterDataSplitter, QbvhProxy};
use crate::simd::{SimdReal, SIMD_WIDTH};
use simba::simd::{SimdBool, SimdValue};
use super::{IndexedData, NodeIndex, Qbvh, QbvhNode, QbvhNodeFlags};
#[cfg(test)]
mod tests;
#[derive(Clone, Default)]
pub struct QbvhUpdateWorkspace {
stack: Vec<(u32, u8)>,
dirty_parent_nodes: Vec<u32>,
to_sort: Vec<usize>,
orig_ids: Vec<u32>,
is_leaf: Vec<bool>,
aabbs: Vec<Aabb>,
}
impl QbvhUpdateWorkspace {
fn clear(&mut self) {
self.stack.clear();
self.dirty_parent_nodes.clear();
self.to_sort.clear();
self.orig_ids.clear();
self.is_leaf.clear();
self.aabbs.clear();
}
}
impl<LeafData: IndexedData> Qbvh<LeafData> {
pub fn remove(&mut self, data: LeafData) -> Option<LeafData> {
let id = data.index();
let proxy = self.proxies.get_mut(id)?;
let node = self.nodes.get_mut(proxy.node.index as usize)?;
node.children[proxy.node.lane as usize] = u32::MAX;
if !node.is_dirty() {
node.set_dirty(true);
self.dirty_nodes.push(proxy.node.index);
}
*proxy = QbvhProxy::invalid();
Some(data)
}
pub fn pre_update_or_insert(&mut self, data: LeafData) {
if self.nodes.is_empty() {
let mut root = QbvhNode::empty();
root.children = [1, u32::MAX, u32::MAX, u32::MAX];
self.nodes
.extend_from_slice(&[root, QbvhNode::empty_leaf_with_parent(NodeIndex::new(0, 0))]);
}
let proxy_id = data.index();
if self.proxies.len() <= proxy_id {
self.proxies.resize(proxy_id + 1, QbvhProxy::invalid());
}
let proxy = &mut self.proxies[proxy_id];
proxy.data = data;
if proxy.is_detached() {
for ii in 0..SIMD_WIDTH {
let mut child = self.nodes[0].children[ii];
if child == u32::MAX {
child = self.nodes.len() as u32;
self.nodes
.push(QbvhNode::empty_leaf_with_parent(NodeIndex::new(
0, ii as u8,
)));
self.nodes[0].children[ii] = child;
}
let child_node = &mut self.nodes[child as usize];
if !child_node.is_leaf() {
continue;
}
for kk in 0..SIMD_WIDTH {
if child_node.children[kk] == u32::MAX {
child_node.children[kk] = proxy_id as u32;
proxy.node = NodeIndex::new(child, kk as u8);
if !child_node.is_dirty() {
self.dirty_nodes.push(child);
child_node.set_dirty(true);
}
return;
}
}
}
let mut old_root = self.nodes[0];
old_root.parent = NodeIndex::new(0, 0);
for child_id in old_root.children {
let parent_index = self.nodes.len() as u32;
if let Some(child_node) = self.nodes.get_mut(child_id as usize) {
child_node.parent.index = parent_index;
}
}
let new_leaf_node_id = self.nodes.len() as u32 + 1;
let mut new_leaf_node = QbvhNode::empty_leaf_with_parent(NodeIndex::new(0, 1));
new_leaf_node.children[0] = proxy_id as u32;
self.dirty_nodes.push(new_leaf_node_id);
new_leaf_node.set_dirty(true);
proxy.node = NodeIndex::new(new_leaf_node_id, 0);
let new_root_children = [
self.nodes.len() as u32,
new_leaf_node_id,
u32::MAX,
u32::MAX,
];
self.nodes.extend_from_slice(&[old_root, new_leaf_node]);
self.nodes[0].children = new_root_children;
} else {
let node = &mut self.nodes[proxy.node.index as usize];
if !node.is_dirty() {
node.set_dirty(true);
self.dirty_nodes.push(proxy.node.index);
}
}
}
#[allow(dead_code)] pub(crate) fn clear_changed_flag(&mut self, workspace: &mut QbvhUpdateWorkspace) {
if self.nodes.is_empty() {
return;
}
workspace.clear();
workspace.stack.push((0, 0));
while let Some((id, _)) = workspace.stack.pop() {
let node = &mut self.nodes[id as usize];
if node.is_changed() {
node.set_changed(false);
for child in node.children {
if child < self.nodes.len() as u32 {
workspace.stack.push((child, 0));
}
}
}
}
}
#[doc(hidden)]
pub fn check_topology(&self, check_aabbs: bool, aabb_builder: impl Fn(&LeafData) -> Aabb) {
if self.nodes.is_empty() {
return;
}
let mut stack = vec![];
stack.push(0);
let mut node_id_found = vec![false; self.nodes.len()];
let mut proxy_id_found = vec![false; self.proxies.len()];
while let Some(id) = stack.pop() {
let node = &self.nodes[id as usize];
assert!(!node_id_found[id as usize]);
node_id_found[id as usize] = true;
if id != 0 {
let parent = &self.nodes[node.parent.index as usize];
assert!(!parent.is_leaf());
assert_eq!(parent.children[node.parent.lane as usize], id);
if check_aabbs {
assert!(
parent
.simd_aabb
.extract(node.parent.lane as usize)
.contains(&node.simd_aabb.to_merged_aabb()),
"failed for {id}"
);
}
if node.is_changed() {
assert!(parent.is_changed());
}
}
if node.is_leaf() {
for ii in 0..SIMD_WIDTH {
let proxy_id = node.children[ii];
if proxy_id != u32::MAX {
assert!(!proxy_id_found[proxy_id as usize]);
proxy_id_found[proxy_id as usize] = true;
if check_aabbs {
let aabb = node.simd_aabb.extract(ii);
assert!(
aabb.contains(&aabb_builder(&self.proxies[proxy_id as usize].data))
);
}
assert_eq!(
self.proxies[proxy_id as usize].node,
NodeIndex::new(id, ii as u8),
"Failed {proxy_id}"
);
}
}
} else {
for child in node.children {
if child != u32::MAX {
stack.push(child);
}
}
}
}
assert_eq!(
proxy_id_found.iter().filter(|found| **found).count(),
self.proxies.iter().filter(|p| !p.is_detached()).count()
);
}
pub fn refit<F>(
&mut self,
margin: Real,
workspace: &mut QbvhUpdateWorkspace,
aabb_builder: F,
) -> usize
where
F: Fn(&LeafData) -> Aabb,
{
workspace.clear();
let margin = SimdReal::splat(margin);
let mut first_iter = true;
let mut num_changed = 0;
while !self.dirty_nodes.is_empty() {
while let Some(id) = self.dirty_nodes.pop() {
if let Some(node) = self.nodes.get(id as usize) {
let mut new_aabbs = [Aabb::new_invalid(); SIMD_WIDTH];
for (child_id, new_aabb) in node.children.iter().zip(new_aabbs.iter_mut()) {
if node.is_leaf() {
if let Some(proxy) = self.proxies.get(*child_id as usize) {
*new_aabb = aabb_builder(&proxy.data);
}
} else {
if let Some(node) = self.nodes.get(*child_id as usize) {
*new_aabb = node.simd_aabb.to_merged_aabb();
}
}
}
let node = &mut self.nodes[id as usize];
let new_simd_aabb = SimdAabb::from(new_aabbs);
node.set_dirty(false);
if !first_iter || !node.simd_aabb.contains(&new_simd_aabb).all() {
node.set_changed(true);
node.simd_aabb = new_simd_aabb;
node.simd_aabb.loosen(margin);
let parent_id = node.parent.index;
num_changed += 1;
if let Some(parent) = self.nodes.get_mut(parent_id as usize) {
if !parent.is_dirty() {
workspace.dirty_parent_nodes.push(parent_id);
parent.set_dirty(true);
}
}
}
}
}
first_iter = false;
std::mem::swap(&mut self.dirty_nodes, &mut workspace.dirty_parent_nodes);
}
num_changed
}
pub fn rebalance(&mut self, margin: Real, workspace: &mut QbvhUpdateWorkspace) {
if self.nodes.is_empty() {
return;
}
workspace.clear();
const MIN_CHANGED_DEPTH: u8 = 5; const FULL_REBUILD_DEPTH: u8 = 15;
for root_child in self.nodes[0].children {
if root_child < self.nodes.len() as u32 {
workspace.stack.push((root_child, 1));
}
}
let mut force_full_rebuild = false;
while let Some((id, depth)) = workspace.stack.pop() {
if depth > FULL_REBUILD_DEPTH {
force_full_rebuild = true;
break;
}
let node = &mut self.nodes[id as usize];
if node.is_leaf() {
self.free_list.push(id);
for ii in 0..SIMD_WIDTH {
let proxy_id = node.children[ii];
if proxy_id < self.proxies.len() as u32 {
workspace.orig_ids.push(proxy_id);
workspace.aabbs.push(node.simd_aabb.extract(ii));
workspace.is_leaf.push(true);
}
}
} else if node.is_changed() || depth < MIN_CHANGED_DEPTH {
self.free_list.push(id);
for child in node.children {
if child < self.nodes.len() as u32 {
workspace.stack.push((child, depth + 1));
}
}
} else {
workspace.orig_ids.push(id);
workspace.aabbs.push(node.simd_aabb.to_merged_aabb());
workspace.is_leaf.push(false);
}
}
if force_full_rebuild {
workspace.clear();
let all_leaves: Vec<_> = self
.proxies
.iter()
.filter_map(|proxy| self.node_aabb(proxy.node).map(|aabb| (proxy.data, aabb)))
.collect();
self.clear_and_rebuild(all_leaves.into_iter(), 0.0);
return;
}
workspace.to_sort.extend(0..workspace.orig_ids.len());
let root_id = NodeIndex::new(0, 0);
let mut indices = std::mem::take(&mut workspace.to_sort);
let (id, aabb) = self.do_recurse_rebalance(&mut indices, workspace, root_id, margin);
workspace.to_sort = indices;
self.root_aabb = aabb;
self.nodes[0] = QbvhNode {
simd_aabb: SimdAabb::from([
aabb,
Aabb::new_invalid(),
Aabb::new_invalid(),
Aabb::new_invalid(),
]),
children: [id, u32::MAX, u32::MAX, u32::MAX],
parent: NodeIndex::invalid(),
flags: QbvhNodeFlags::default(),
};
}
fn do_recurse_rebalance(
&mut self,
indices: &mut [usize],
workspace: &QbvhUpdateWorkspace,
parent: NodeIndex,
margin: Real,
) -> (u32, Aabb) {
if indices.len() <= 4 {
let mut has_leaf = false;
let mut has_internal = false;
for id in &*indices {
has_leaf |= workspace.is_leaf[*id];
has_internal |= !workspace.is_leaf[*id];
}
let my_internal_id = if has_internal {
self.free_list.pop().unwrap_or_else(|| {
self.nodes.push(QbvhNode::empty());
self.nodes.len() as u32 - 1
})
} else {
u32::MAX
};
let my_leaf_id = if has_leaf {
self.free_list.pop().unwrap_or_else(|| {
self.nodes.push(QbvhNode::empty());
self.nodes.len() as u32 - 1
})
} else {
u32::MAX
};
let mut my_leaf_aabb = Aabb::new_invalid();
let mut my_internal_aabb = Aabb::new_invalid();
let mut leaf_aabbs = [Aabb::new_invalid(); 4];
let mut internal_aabbs = [Aabb::new_invalid(); 4];
let mut proxy_ids = [u32::MAX; 4];
let mut internal_ids = [u32::MAX; 4];
let mut new_internal_lane_containing_leaf = usize::MAX;
for (k, id) in indices.iter().enumerate() {
let orig_id = workspace.orig_ids[*id];
if workspace.is_leaf[*id] {
new_internal_lane_containing_leaf = k;
my_leaf_aabb.merge(&workspace.aabbs[*id]);
leaf_aabbs[k] = workspace.aabbs[*id];
proxy_ids[k] = orig_id;
self.proxies[orig_id as usize].node = NodeIndex::new(my_leaf_id, k as u8);
} else {
my_internal_aabb.merge(&workspace.aabbs[*id]);
internal_aabbs[k] = workspace.aabbs[*id];
internal_ids[k] = orig_id;
self.nodes[orig_id as usize].parent = NodeIndex::new(my_internal_id, k as u8);
}
}
if has_internal {
if has_leaf {
internal_ids[new_internal_lane_containing_leaf] = my_leaf_id;
internal_aabbs[new_internal_lane_containing_leaf] = my_leaf_aabb;
my_internal_aabb.merge(&my_leaf_aabb);
}
let internal_node = QbvhNode {
simd_aabb: SimdAabb::from(internal_aabbs),
children: internal_ids,
parent,
flags: QbvhNodeFlags::default(),
};
self.nodes[my_internal_id as usize] = internal_node;
}
if has_leaf {
let leaf_node = QbvhNode {
simd_aabb: SimdAabb::from(leaf_aabbs),
children: proxy_ids,
parent: if has_internal {
NodeIndex::new(my_internal_id, new_internal_lane_containing_leaf as u8)
} else {
parent
},
flags: QbvhNodeFlags::LEAF,
};
self.nodes[my_leaf_id as usize] = leaf_node;
}
if has_internal {
return (my_internal_id, my_internal_aabb);
} else {
return (my_leaf_id, my_leaf_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 = workspace.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 = workspace.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 nid = if let Some(nid) = self.free_list.pop() {
self.nodes[nid as usize] = node;
nid
} else {
let nid = self.nodes.len();
self.nodes.push(node);
nid as u32
};
let splitter = CenterDataSplitter::default();
let splits =
splitter.split_dataset_wo_workspace(subdiv_dims, center, indices, &workspace.aabbs);
let n = [
NodeIndex::new(nid, 0),
NodeIndex::new(nid, 1),
NodeIndex::new(nid, 2),
NodeIndex::new(nid, 3),
];
let children = [
self.do_recurse_rebalance(splits[0], workspace, n[0], margin),
self.do_recurse_rebalance(splits[1], workspace, n[1], margin),
self.do_recurse_rebalance(splits[2], workspace, n[2], margin),
self.do_recurse_rebalance(splits[3], workspace, n[3], margin),
];
self.nodes[nid as usize].children =
[children[0].0, children[1].0, children[2].0, children[3].0];
self.nodes[nid as usize].simd_aabb =
SimdAabb::from([children[0].1, children[1].1, children[2].1, children[3].1]);
self.nodes[nid as usize]
.simd_aabb
.loosen(SimdReal::splat(margin));
let my_aabb = self.nodes[nid as usize].simd_aabb.to_merged_aabb();
(nid, my_aabb)
}
}