nodedb_spatial/rtree/
node.rs

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

nodedb_spatial/rtree/
node.rs
