use std::mem;
use pi_link_list::{LinkList, Node};
use pi_null::Null;
use pi_slotmap::{new_key_type, Key, SecondaryMap, SlotMap};
new_key_type! {
pub struct BranchKey;
}
pub trait Helper<const N: usize> {
type Point;
type Vector: Clone;
type Aabb: Clone;
fn aabb_extents(aabb: &Self::Aabb) -> Self::Vector;
fn aabb_shift(aabb: &Self::Aabb, distance: &Self::Vector) -> Self::Aabb;
fn aabb_contains(aabb: &Self::Aabb, other: &Self::Aabb) -> bool;
fn aabb_intersects(aabb: &Self::Aabb, other: &Self::Aabb) -> bool;
fn get_deap(
d: &mut Self::Vector,
loose_layer: usize,
max_loose: &Self::Vector,
deep: usize,
min_loose: &Self::Vector,
) -> usize;
fn smaller_than_min_loose(d: &Self::Vector, min_loose: &Self::Vector) -> bool;
fn calc_layer(loose: &Self::Vector, el: &Self::Vector) -> usize;
fn get_child(point: &Self::Point, aabb: &Self::Aabb) -> u8;
fn get_max_half_loose(aabb: &Self::Aabb, loose: &Self::Vector) -> Self::Point;
fn make_childs(aabb: &Self::Aabb, loose: &Self::Vector) -> [Self::Aabb; N];
fn create_child(
aabb: &Self::Aabb,
loose: &Self::Vector,
layer: usize,
loose_layer: usize,
min_loose: &Self::Vector,
child_index: u8,
) -> (Self::Aabb, Self::Vector);
}
const DEEP_MAX: usize = 16;
const ADJUST_MIN: usize = 4;
const ADJUST_MAX: usize = 8;
const AUTO_COLLECT: usize = 1024;
type List<K, H, T, const N: usize> = LinkList<
K,
AbNode<<H as Helper<N>>::Aabb, T>,
SecondaryMap<K, Node<K, AbNode<<H as Helper<N>>::Aabb, T>>>,
>;
pub struct Tree<K: Key, H: Helper<N>, T, const N: usize> {
pub slab: SlotMap<BranchKey, BranchNode<K, H, T, N>>, pub ab_map: SecondaryMap<K, Node<K, AbNode<H::Aabb, T>>>, max_loose: H::Vector, min_loose: H::Vector, root_key: BranchKey,
pub outer: List<K, H, T, N>, pub dirty: (Vec<Vec<BranchKey>>, DirtyState), adjust: (usize, usize), loose_layer: usize, deep: usize, auto_collect: usize, }
impl<K: Key, H: Helper<N>, T, const N: usize> Tree<K, H, T, N> {
pub fn new(
root: H::Aabb,
max_loose: H::Vector,
min_loose: H::Vector,
adjust_min: usize,
adjust_max: usize,
deep: usize,
) -> Self {
let adjust_min = if adjust_min == 0 {
ADJUST_MIN
} else {
adjust_min
};
let adjust_max = if adjust_max == 0 {
ADJUST_MAX
} else {
adjust_max
};
let adjust_max = adjust_min.max(adjust_max);
let deep = if deep > DEEP_MAX || deep == 0 {
DEEP_MAX
} else {
deep
};
let mut branch_slab: SlotMap<BranchKey, BranchNode<K, H, T, N>> = SlotMap::with_key();
let mut d = H::aabb_extents(&root);
let loose_layer = H::calc_layer(&max_loose, &min_loose);
let deep = H::get_deap(&mut d, loose_layer, &max_loose, deep, &min_loose);
let root = branch_slab.insert(BranchNode::new(
root,
max_loose.clone(),
0,
BranchKey::null(),
0,
));
return Tree {
slab: branch_slab,
ab_map: SecondaryMap::default(),
max_loose,
min_loose,
adjust: (adjust_min, adjust_max),
loose_layer,
deep,
root_key: root,
outer: LinkList::new(),
dirty: (
Vec::new(),
DirtyState {
dirty_count: 0,
min_layer: usize::max_value(),
max_layer: 0,
},
),
auto_collect: AUTO_COLLECT,
};
}
pub fn get_auto_collect(&self) -> usize {
self.auto_collect
}
pub fn set_auto_collect(&mut self, auto_collect: usize) {
self.auto_collect = auto_collect;
}
pub fn get_adjust(&self) -> (usize, usize) {
(self.adjust.0, self.adjust.1)
}
pub fn get_layer(&self, aabb: &H::Aabb) -> usize {
let d = H::aabb_extents(aabb);
if H::smaller_than_min_loose(&d, &self.min_loose) {
return self.deep;
};
H::calc_layer(&self.max_loose, &d)
}
pub fn add(&mut self, id: K, aabb: H::Aabb, bind: T) -> bool {
if self.ab_map.contains_key(id) {
return false;
}
let layer = self.get_layer(&aabb);
self.ab_map.insert(
id,
Node::new(AbNode::new(aabb.clone(), bind, layer, N as u8)),
);
let root = unsafe { self.slab.get_unchecked_mut(self.root_key) };
if H::aabb_contains(&root.aabb, &aabb) {
self.down(self.root_key, &aabb, layer, id);
} else {
self.outer.link_before(id, K::null(), &mut self.ab_map);
}
true
}
fn down(&mut self, branch_id: BranchKey, aabb: &H::Aabb, layer: usize, id: K) {
let parent = unsafe { self.slab.get_unchecked_mut(branch_id) };
let child = if parent.layer as usize >= layer {
parent.nodes.link_before(id, K::null(), &mut self.ab_map);
N as u8
} else {
let i = H::get_child(&H::get_max_half_loose(&parent.aabb, &parent.loose), aabb);
match parent.childs[i as usize] {
ChildNode::Branch(branch) => {
return self.down(branch, aabb, layer, id);
}
ChildNode::Ab(ref mut list) => {
list.link_before(id, K::null(), &mut self.ab_map);
if list.len() >= self.adjust.1 && parent.layer < self.deep {
set_dirty(&mut parent.dirty, parent.layer, branch_id, &mut self.dirty);
}
}
}
i
};
let node = unsafe { self.ab_map.get_unchecked_mut(id) };
node.parent = branch_id;
node.parent_child = child;
if self.dirty.1.dirty_count >= self.auto_collect {
self.collect();
}
}
pub fn get(&self, id: K) -> Option<&(H::Aabb, T)> {
match self.ab_map.get(id) {
Some(node) => Some(&node.value),
_ => None,
}
}
pub unsafe fn get_unchecked(&self, id: K) -> &(H::Aabb, T) {
&self.ab_map.get_unchecked(id).value
}
pub unsafe fn get_mut(&mut self, id: K) -> Option<&mut T> {
match self.ab_map.get_mut(id) {
Some(n) => Some(&mut n.value.1),
_ => None,
}
}
pub unsafe fn get_unchecked_mut(&mut self, id: K) -> &mut T {
let node = self.ab_map.get_unchecked_mut(id);
&mut node.value.1
}
pub fn contains_key(&self, id: K) -> bool {
self.ab_map.contains_key(id)
}
pub fn update(&mut self, id: K, aabb: H::Aabb) -> bool {
let layer = self.get_layer(&aabb);
if let Some(node) = self.ab_map.get_mut(id) {
node.layer = layer;
node.value.0 = aabb.clone();
let old_p = node.parent;
let old_c = node.parent_child;
self.update1(id, layer, old_p, old_c, &aabb);
true
} else {
false
}
}
fn update1(&mut self, id: K, layer: usize, old_p: BranchKey, old_c: u8, aabb: &H::Aabb) {
if old_p.is_null() {
let root = unsafe { self.slab.get_unchecked_mut(self.root_key) };
if H::aabb_contains(&root.aabb, aabb) {
self.outer.unlink(id, &mut self.ab_map);
self.down(self.root_key, aabb, layer, id);
} else {
}
return;
}
let mut parent = unsafe { self.slab.get_unchecked_mut(old_p) };
if layer > parent.layer {
if H::aabb_contains(&parent.aabb, aabb) {
let child = H::get_child(&H::get_max_half_loose(&parent.aabb, &parent.loose), aabb);
if old_c == child {
return;
}
Self::remove1(&mut self.ab_map, id, old_c, parent);
match parent.childs[child as usize] {
ChildNode::Branch(branch) => {
self.down(branch, aabb, layer, id);
}
ChildNode::Ab(ref mut list) => {
Self::add1(&mut self.ab_map, list, id, old_p, child);
if list.len() >= self.adjust.1 && layer < self.deep {
set_dirty(&mut parent.dirty, parent.layer, old_p, &mut self.dirty);
}
}
}
return;
}
} else if layer == parent.layer {
if H::aabb_contains(&parent.aabb, aabb) {
if (old_c as usize) == N {
return;
}
Self::remove1(&mut self.ab_map, id, old_c, parent);
Self::add1(&mut self.ab_map, &mut parent.nodes, id, old_p, N as u8);
return;
}
} else {
};
Self::remove1(&mut self.ab_map, id, old_c, parent);
if parent.is_need_merge(self.adjust.0) {
set_dirty(&mut parent.dirty, parent.layer, old_p, &mut self.dirty);
}
let mut p = parent.parent;
while !p.is_null() {
parent = unsafe { self.slab.get_unchecked_mut(p) };
if parent.layer <= layer && H::aabb_contains(&parent.aabb, aabb) {
return self.down(p, aabb, layer, id);
}
p = parent.parent;
}
Self::add1(
&mut self.ab_map,
&mut self.outer,
id,
BranchKey::null(),
N as u8,
);
}
fn remove1(
ab_map: &mut SecondaryMap<K, Node<K, AbNode<H::Aabb, T>>>,
id: K,
old_c: u8,
parent: &mut BranchNode<K, H, T, N>,
) {
if (old_c as usize) < N {
match parent.childs[old_c as usize] {
ChildNode::Ab(ref mut list) => list.unlink(id, ab_map),
_ => panic!("invalid state"),
}
} else {
parent.nodes.unlink(id, ab_map);
}
}
fn add1(
ab_map: &mut SecondaryMap<K, Node<K, AbNode<H::Aabb, T>>>,
list: &mut List<K, H, T, N>,
id: K,
parent: BranchKey,
parent_child: u8,
) {
let node = unsafe { ab_map.get_unchecked_mut(id) };
node.parent = parent;
node.parent_child = parent_child;
list.link_before(id, K::null(), ab_map);
}
pub fn shift(&mut self, id: K, distance: H::Vector) -> bool {
if let Some(node) = self.ab_map.get_mut(id) {
let aabb = H::aabb_shift(&node.value.0, &distance);
let layer = node.layer;
node.value.0 = aabb.clone();
let old_p = node.parent;
let old_c = node.parent_child;
self.update1(id, layer, old_p, old_c, &aabb);
true
} else {
false
}
}
pub fn update_bind(&mut self, id: K, bind: T) -> bool {
match self.ab_map.get_mut(id) {
Some(node) => {
node.value.1 = bind;
true
}
_ => false,
}
}
pub fn remove(&mut self, id: K) -> Option<(H::Aabb, T)> {
let (parent, parent_child) = match self.ab_map.get(id) {
Some(n) => (n.parent, n.parent_child),
_ => return None,
};
if !parent.is_null() {
let branch = unsafe { self.slab.get_unchecked_mut(parent) };
Self::remove1(&mut self.ab_map, id, parent_child, branch);
if branch.is_need_merge(self.adjust.0) {
set_dirty(&mut branch.dirty, branch.layer, parent, &mut self.dirty);
}
} else {
self.outer.unlink(id, &mut self.ab_map);
}
Some(self.ab_map.remove(id).unwrap().take().value)
}
pub fn collect(&mut self) {
let state = mem::replace(&mut self.dirty.1, DirtyState::new());
if state.dirty_count == 0 {
return;
}
for i in state.min_layer..state.max_layer {
let vec = unsafe { self.dirty.0.get_unchecked_mut(i) };
let c = vec.len();
if c == 0 {
continue;
}
for j in 0..c {
let branch_id = unsafe { vec.get_unchecked(j) };
Self::collect1(
&mut self.slab,
&mut self.ab_map,
&self.adjust,
self.deep,
*branch_id,
self.loose_layer,
&self.min_loose,
);
}
vec.clear();
}
}
fn collect1(
slab: &mut SlotMap<BranchKey, BranchNode<K, H, T, N>>,
ab_map: &mut SecondaryMap<K, Node<K, AbNode<H::Aabb, T>>>,
adjust: &(usize, usize),
deep: usize,
branch_id: BranchKey,
loose_layer: usize,
min_loose: &H::Vector,
) {
let parent = match slab.get_mut(branch_id) {
Some(branch) => branch,
_ => return,
};
let dirty = mem::replace(&mut parent.dirty, false);
if !dirty {
return;
}
let parent_id = parent.parent;
if (!parent_id.is_null()) && parent.is_need_merge(adjust.0) {
let child = parent.parent_child;
let list = Self::merge_branch(ab_map, parent, LinkList::new());
slab.remove(branch_id);
Self::shrink(slab, ab_map, adjust.0, parent_id, child, branch_id, list);
return;
}
let (need, lists) = parent.need_split_list(adjust.1);
if need {
let aabb = parent.aabb.clone();
let loose = parent.loose.clone();
let layer = parent.layer;
Self::split(
slab,
ab_map,
adjust.1,
deep,
lists,
&aabb,
&loose,
layer,
branch_id,
loose_layer,
min_loose,
);
}
}
fn merge_branch(
ab_map: &mut SecondaryMap<K, Node<K, AbNode<H::Aabb, T>>>,
branch: &mut BranchNode<K, H, T, N>,
mut list: List<K, H, T, N>,
) -> List<K, H, T, N> {
list.append(&mut branch.nodes, ab_map);
for n in &mut branch.childs {
match n {
ChildNode::Ab(other) => list.append(other, ab_map),
_ => (),
}
}
list
}
fn shrink(
slab: &mut SlotMap<BranchKey, BranchNode<K, H, T, N>>,
ab_map: &mut SecondaryMap<K, Node<K, AbNode<H::Aabb, T>>>,
adjust: usize,
branch_id: BranchKey,
parent_child: u8,
child_id: BranchKey,
list: List<K, H, T, N>,
) {
let branch = unsafe { slab.get_unchecked_mut(branch_id) };
if (!branch.parent.is_null()) && branch.is_need_merge_with_child(adjust, child_id, list.len()) {
let parent_id = branch.parent;
let child = branch.parent_child;
let list = Self::merge_branch(ab_map, branch, list);
slab.remove(branch_id);
Self::shrink(slab, ab_map, adjust, parent_id, child, branch_id, list);
} else {
for (_, node) in list.iter_mut(ab_map) {
node.parent = branch_id;
node.parent_child = parent_child;
};
branch.childs[parent_child as usize] = ChildNode::Ab(list);
}
}
#[inline]
fn split(
slab: &mut SlotMap<BranchKey, BranchNode<K, H, T, N>>,
ab_map: &mut SecondaryMap<K, Node<K, AbNode<H::Aabb, T>>>,
adjust: usize,
deep: usize,
lists: [List<K, H, T, N>; N],
parent_aabb: &H::Aabb,
parent_loose: &H::Vector,
parent_layer: usize,
parent_id: BranchKey,
loose_layer: usize,
min_loose: &H::Vector,
) {
let mut branchs = [BranchKey::null(); N];
for (i, list) in lists.into_iter().enumerate() {
if list.is_empty() {
continue;
}
let branch = BranchNode::create(
parent_aabb,
parent_loose,
parent_layer,
parent_id,
loose_layer,
min_loose,
i as u8,
);
let branch_id = slab.insert(branch);
Self::split_down(
slab,
ab_map,
adjust,
deep,
list,
branch_id,
loose_layer,
min_loose,
);
branchs[i] = branch_id;
}
let parent = unsafe { slab.get_unchecked_mut(parent_id) };
for (i, child_id) in branchs.into_iter().enumerate() {
if !child_id.is_null() {
parent.childs[i] = ChildNode::Branch(child_id);
}
}
}
fn split_down(
slab: &mut SlotMap<BranchKey, BranchNode<K, H, T, N>>,
ab_map: &mut SecondaryMap<K, Node<K, AbNode<H::Aabb, T>>>,
adjust: usize,
deep: usize,
list: List<K, H, T, N>,
parent_id: BranchKey,
loose_layer: usize,
min_loose: &H::Vector,
) {
let parent = unsafe { slab.get_unchecked_mut(parent_id) };
let point = H::get_max_half_loose(&parent.aabb, &parent.loose);
let mut drain = list.drain();
let mut id = drain.pop_front(ab_map);
while !id.is_null() {
let node = unsafe { ab_map.get_unchecked_mut(id) };
if parent.layer >= node.layer {
node.parent = parent_id;
node.parent_child = N as u8;
parent.nodes.link_before(id, K::null(), ab_map);
} else {
let i = H::get_child(&point, &node.value.0);
match parent.childs[i as usize] {
ChildNode::Ab(ref mut ab) => {
node.parent = parent_id;
node.parent_child = i;
ab.link_before(id, K::null(), ab_map);
}
_ => panic!("invalid state"),
}
}
id = drain.pop_front(ab_map);
}
if parent.layer >= deep {
return;
}
let (need, lists) = parent.need_split_list(adjust);
if need {
let aabb: <H as Helper<N>>::Aabb = parent.aabb.clone();
let loose = parent.loose.clone();
let layer = parent.layer;
Self::split(
slab,
ab_map,
adjust,
deep,
lists,
&aabb,
&loose,
layer,
parent_id,
loose_layer,
min_loose,
);
}
}
pub fn query<A, B>(
&self,
branch_arg: &A,
branch_func: fn(arg: &A, aabb: &H::Aabb) -> bool,
ab_arg: &mut B,
ab_func: fn(arg: &mut B, id: K, aabb: &H::Aabb, bind: &T),
) {
self.query_outer(ab_arg, ab_func);
self.query1(self.root_key, branch_arg, branch_func, ab_arg, ab_func)
}
fn query1<A, B>(
&self,
branch_id: BranchKey,
branch_arg: &A,
branch_func: fn(arg: &A, aabb: &H::Aabb) -> bool,
ab_arg: &mut B,
ab_func: fn(arg: &mut B, id: K, aabb: &H::Aabb, bind: &T),
) {
let node = unsafe { self.slab.get_unchecked(branch_id) };
for (id, ab) in node.nodes.iter(&self.ab_map) {
ab_func(ab_arg, id, &ab.value.0, &ab.value.1);
}
let childs = H::make_childs(&node.aabb, &node.loose);
for (i, ab) in childs.iter().enumerate() {
match node.childs[i] {
ChildNode::Branch(branch) => {
if branch_func(branch_arg, &ab) {
self.query1(branch, branch_arg, branch_func, ab_arg, ab_func);
}
}
ChildNode::Ab(ref list) if !list.is_empty() => {
if branch_func(branch_arg, &ab) {
for (id, ab) in list.iter(&self.ab_map) {
ab_func(ab_arg, id, &ab.value.0, &ab.value.1);
}
}
}
_ => (),
}
}
}
pub fn query_outer<B>(
&self,
arg: &mut B,
func: fn(arg: &mut B, id: K, aabb: &H::Aabb, bind: &T),
) {
for (id, ab) in self.outer.iter(&self.ab_map) {
func(arg, id, &ab.value.0, &ab.value.1);
}
}
pub fn len(&self) -> usize {
self.ab_map.len()
}
}
#[derive(Clone)]
pub struct BranchNode<K: Key, H: Helper<N>, T, const N: usize> {
aabb: H::Aabb, loose: H::Vector, layer: usize, parent: BranchKey, childs: [ChildNode<K, H, T, N>; N], nodes: List<K, H, T, N>, parent_child: u8, dirty: bool, }
impl<K: Key, H: Helper<N>, T, const N: usize> BranchNode<K, H, T, N> {
#[inline]
pub fn new(
aabb: H::Aabb,
loose: H::Vector,
layer: usize,
parent: BranchKey,
child: u8,
) -> Self {
let childs = [0; N].map(|_| ChildNode::Ab(Default::default()));
BranchNode {
aabb,
loose,
layer,
parent,
childs,
nodes: LinkList::new(),
parent_child: child,
dirty: false,
}
}
fn create(
aabb: &H::Aabb,
loose: &H::Vector,
layer: usize,
parent_id: BranchKey,
loose_layer: usize,
min_loose: &H::Vector,
child: u8,
) -> Self {
let (ab, loose) = H::create_child(aabb, loose, layer, loose_layer, min_loose, child);
BranchNode::new(ab, loose, layer + 1, parent_id, child)
}
pub fn is_need_merge(&self, adjust_min: usize) -> bool {
if self.parent.is_null() {
return false;
}
let mut len = self.nodes.len();
for n in &self.childs {
match n {
ChildNode::Branch(_) => return false,
ChildNode::Ab(list) => len += list.len(),
}
}
len <= adjust_min
}
pub fn is_need_merge_with_child(
&self,
adjust_min: usize,
child: BranchKey,
child_node_len: usize,
) -> bool {
let mut len = self.nodes.len();
for n in &self.childs {
match n {
ChildNode::Branch(b) => {
if b != &child {
return false;
}
len += child_node_len;
}
ChildNode::Ab(list) => len += list.len(),
}
}
len <= adjust_min
}
pub fn need_split_list(&mut self, adjust_max: usize) -> (bool, [List<K, H, T, N>; N]) {
let mut need = false;
let mut childs = [0; N].map(|_| Default::default());
for (i, n) in self.childs.iter_mut().enumerate() {
match n {
ChildNode::Ab(list) if list.len() >= adjust_max => {
mem::swap(list, &mut childs[i]);
need = true;
}
_ => (),
}
}
(need, childs)
}
}
#[derive(Clone)]
enum ChildNode<K: Key, H: Helper<N>, T, const N: usize> {
Branch(BranchKey), Ab(List<K, H, T, N>), }
#[derive(Debug, Clone)]
pub struct AbNode<Aabb, T> {
value: (Aabb, T), parent: BranchKey, layer: usize, parent_child: u8, }
impl<Aabb, T> AbNode<Aabb, T> {
pub fn new(aabb: Aabb, bind: T, layer: usize, n: u8) -> Self {
AbNode {
value: (aabb, bind),
layer: layer,
parent: BranchKey::null(),
parent_child: n,
}
}
}
#[derive(Debug)]
pub struct DirtyState {
dirty_count: usize,
min_layer: usize,
max_layer: usize,
}
impl DirtyState {
fn new() -> Self {
DirtyState {
dirty_count: 0,
min_layer: usize::max_value(),
max_layer: 0,
}
}
}
#[inline]
fn set_dirty(
dirty: &mut bool,
layer: usize,
rid: BranchKey,
dirty_list: &mut (Vec<Vec<BranchKey>>, DirtyState),
) {
dirty_list.1.dirty_count += 1;
if !*dirty {
set_tree_dirty(dirty_list, layer, rid);
}
*dirty = true;
}
#[inline]
fn set_tree_dirty(dirty: &mut (Vec<Vec<BranchKey>>, DirtyState), layer: usize, rid: BranchKey) {
if dirty.1.min_layer > layer {
dirty.1.min_layer = layer;
}
if dirty.1.max_layer <= layer {
dirty.1.max_layer = layer + 1;
}
if dirty.0.len() <= layer as usize {
for _ in dirty.0.len()..layer as usize + 1 {
dirty.0.push(Vec::new())
}
}
let vec = unsafe { dirty.0.get_unchecked_mut(layer as usize) };
vec.push(rid);
}