1use alloc::vec::Vec;
8
9use crate::bounds::{Bounds, union_all};
10use crate::indexable::Indexable;
11
12#[derive(Debug, Clone)]
18pub enum Node<T> {
19 Leaf(Vec<T>),
21 Branch(Vec<(Bounds, Node<T>)>),
23}
24
25impl<T: Indexable> Node<T> {
26 #[must_use]
29 pub fn bounds(&self) -> Option<Bounds> {
30 match self {
31 Node::Leaf(values) => {
32 if values.is_empty() {
33 None
34 } else {
35 let boxes: Vec<Bounds> = values.iter().map(Indexable::bounds).collect();
36 Some(union_all(&boxes))
37 }
38 }
39 Node::Branch(children) => {
40 if children.is_empty() {
41 None
42 } else {
43 let boxes: Vec<Bounds> = children.iter().map(|(b, _)| *b).collect();
44 Some(union_all(&boxes))
45 }
46 }
47 }
48 }
49
50 #[must_use]
53 pub fn entry_count(&self) -> usize {
54 match self {
55 Node::Leaf(values) => values.len(),
56 Node::Branch(children) => children.len(),
57 }
58 }
59
60 #[must_use]
62 pub fn value_count(&self) -> usize {
63 match self {
64 Node::Leaf(values) => values.len(),
65 Node::Branch(children) => children.iter().map(|(_, child)| child.value_count()).sum(),
66 }
67 }
68
69 #[must_use]
71 pub fn height(&self) -> usize {
72 match self {
73 Node::Leaf(_) => 1,
74 Node::Branch(children) => {
75 1 + children
76 .iter()
77 .map(|(_, child)| child.height())
78 .max()
79 .unwrap_or(0)
80 }
81 }
82 }
83}
84
85#[cfg(test)]
86mod tests {
87 use super::Node;
88 use crate::bounds::Bounds;
89
90 #[test]
91 fn leaf_bounds_union_of_values() {
92 let leaf: Node<Bounds> = Node::Leaf(alloc::vec![
93 Bounds::point([0.0, 0.0]),
94 Bounds::point([2.0, 3.0]),
95 ]);
96 assert_eq!(leaf.bounds(), Some(Bounds::new([0.0, 0.0], [2.0, 3.0])));
97 assert_eq!(leaf.entry_count(), 2);
98 assert_eq!(leaf.height(), 1);
99 }
100
101 #[test]
102 fn branch_bounds_and_height() {
103 let leaf_a: Node<Bounds> = Node::Leaf(alloc::vec![Bounds::point([0.0, 0.0])]);
104 let leaf_b: Node<Bounds> = Node::Leaf(alloc::vec![Bounds::point([5.0, 5.0])]);
105 let branch: Node<Bounds> = Node::Branch(alloc::vec![
106 (leaf_a.bounds().unwrap(), leaf_a),
107 (leaf_b.bounds().unwrap(), leaf_b),
108 ]);
109 assert_eq!(branch.bounds(), Some(Bounds::new([0.0, 0.0], [5.0, 5.0])));
110 assert_eq!(branch.height(), 2);
111 assert_eq!(branch.value_count(), 2);
112 }
113
114 #[test]
115 fn empty_leaf_has_no_bounds() {
116 let leaf: Node<Bounds> = Node::Leaf(alloc::vec![]);
117 assert_eq!(leaf.bounds(), None);
118 }
119
120 #[test]
121 fn empty_branch_has_no_bounds_or_values() {
122 let branch: Node<Bounds> = Node::Branch(alloc::vec![]);
123 assert_eq!(branch.bounds(), None);
124 assert_eq!(branch.entry_count(), 0);
125 assert_eq!(branch.value_count(), 0);
126 assert_eq!(branch.height(), 1);
127 }
128}