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