Skip to main content

nodedb_spatial/rtree/
node.rs

1//! R-tree node types and constants.
2
3use nodedb_types::BoundingBox;
4use serde::{Deserialize, Serialize};
5
6pub const LEAF_CAPACITY: usize = 32;
7pub const INTERNAL_CAPACITY: usize = 16;
8pub const MIN_FILL_LEAF: usize = LEAF_CAPACITY * 2 / 5;
9pub const MIN_FILL_INTERNAL: usize = INTERNAL_CAPACITY * 2 / 5;
10pub const REINSERT_COUNT_LEAF: usize = LEAF_CAPACITY * 3 / 10;
11
12/// Unique identifier for a spatial entry (document RID or opaque u64).
13pub type EntryId = u64;
14
15/// A spatial entry stored in an R-tree leaf.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct RTreeEntry {
18    pub id: EntryId,
19    pub bbox: BoundingBox,
20}
21
22/// Internal node child reference.
23#[derive(Debug, Clone)]
24pub(crate) struct ChildRef {
25    pub bbox: BoundingBox,
26    pub node_idx: usize,
27}
28
29/// A node in the R-tree (either internal or leaf).
30#[derive(Debug, Clone)]
31pub(crate) enum NodeKind {
32    Internal { children: Vec<ChildRef> },
33    Leaf { entries: Vec<RTreeEntry> },
34}
35
36#[derive(Debug, Clone)]
37pub(crate) struct Node {
38    pub kind: NodeKind,
39    pub bbox: BoundingBox,
40    pub level: u32,
41}
42
43impl Node {
44    pub fn new_leaf() -> Self {
45        Self {
46            kind: NodeKind::Leaf {
47                entries: Vec::with_capacity(LEAF_CAPACITY),
48            },
49            bbox: BoundingBox::new(
50                f64::INFINITY,
51                f64::INFINITY,
52                f64::NEG_INFINITY,
53                f64::NEG_INFINITY,
54            ),
55            level: 0,
56        }
57    }
58
59    pub fn new_internal(level: u32) -> Self {
60        Self {
61            kind: NodeKind::Internal {
62                children: Vec::with_capacity(INTERNAL_CAPACITY),
63            },
64            bbox: BoundingBox::new(
65                f64::INFINITY,
66                f64::INFINITY,
67                f64::NEG_INFINITY,
68                f64::NEG_INFINITY,
69            ),
70            level,
71        }
72    }
73
74    pub fn is_leaf(&self) -> bool {
75        matches!(self.kind, NodeKind::Leaf { .. })
76    }
77
78    pub fn capacity(&self) -> usize {
79        if self.is_leaf() {
80            LEAF_CAPACITY
81        } else {
82            INTERNAL_CAPACITY
83        }
84    }
85
86    pub fn len(&self) -> usize {
87        match &self.kind {
88            NodeKind::Internal { children } => children.len(),
89            NodeKind::Leaf { entries } => entries.len(),
90        }
91    }
92
93    pub fn is_overflow(&self) -> bool {
94        self.len() > self.capacity()
95    }
96
97    pub fn min_fill(&self) -> usize {
98        if self.is_leaf() {
99            MIN_FILL_LEAF
100        } else {
101            MIN_FILL_INTERNAL
102        }
103    }
104
105    pub fn is_underflow(&self) -> bool {
106        self.len() < self.min_fill()
107    }
108
109    pub fn recompute_bbox(&mut self) {
110        match &self.kind {
111            NodeKind::Internal { children } => {
112                if children.is_empty() {
113                    self.bbox = BoundingBox::new(0.0, 0.0, 0.0, 0.0);
114                    return;
115                }
116                let mut bb = children[0].bbox;
117                for c in &children[1..] {
118                    bb = bb.union(&c.bbox);
119                }
120                self.bbox = bb;
121            }
122            NodeKind::Leaf { entries } => {
123                if entries.is_empty() {
124                    self.bbox = BoundingBox::new(0.0, 0.0, 0.0, 0.0);
125                    return;
126                }
127                let mut bb = entries[0].bbox;
128                for e in &entries[1..] {
129                    bb = bb.union(&e.bbox);
130                }
131                self.bbox = bb;
132            }
133        }
134    }
135}