graphos_core/index/
btree.rs

1//! BTree index for range queries.
2
3use graphos_common::types::NodeId;
4use parking_lot::RwLock;
5use std::collections::BTreeMap;
6use std::ops::RangeBounds;
7
8/// A concurrent BTree index for range queries.
9///
10/// This index provides O(log n) lookup and efficient range scans.
11pub struct BTreeIndex<K: Ord, V: Copy> {
12    /// The underlying BTree map.
13    map: RwLock<BTreeMap<K, V>>,
14}
15
16impl<K: Ord + Clone, V: Copy> BTreeIndex<K, V> {
17    /// Creates a new empty BTree index.
18    #[must_use]
19    pub fn new() -> Self {
20        Self {
21            map: RwLock::new(BTreeMap::new()),
22        }
23    }
24
25    /// Inserts a key-value pair into the index.
26    ///
27    /// Returns the previous value if the key was already present.
28    pub fn insert(&self, key: K, value: V) -> Option<V> {
29        self.map.write().insert(key, value)
30    }
31
32    /// Gets the value for a key.
33    pub fn get(&self, key: &K) -> Option<V> {
34        self.map.read().get(key).copied()
35    }
36
37    /// Removes a key from the index.
38    ///
39    /// Returns the value if the key was present.
40    pub fn remove(&self, key: &K) -> Option<V> {
41        self.map.write().remove(key)
42    }
43
44    /// Checks if a key exists in the index.
45    pub fn contains(&self, key: &K) -> bool {
46        self.map.read().contains_key(key)
47    }
48
49    /// Returns all values in the given range.
50    pub fn range<R: RangeBounds<K>>(&self, range: R) -> Vec<(K, V)> {
51        self.map
52            .read()
53            .range(range)
54            .map(|(k, v)| (k.clone(), *v))
55            .collect()
56    }
57
58    /// Returns the minimum key-value pair.
59    pub fn min(&self) -> Option<(K, V)> {
60        self.map
61            .read()
62            .first_key_value()
63            .map(|(k, v)| (k.clone(), *v))
64    }
65
66    /// Returns the maximum key-value pair.
67    pub fn max(&self) -> Option<(K, V)> {
68        self.map
69            .read()
70            .last_key_value()
71            .map(|(k, v)| (k.clone(), *v))
72    }
73
74    /// Returns the number of entries in the index.
75    pub fn len(&self) -> usize {
76        self.map.read().len()
77    }
78
79    /// Returns true if the index is empty.
80    pub fn is_empty(&self) -> bool {
81        self.map.read().is_empty()
82    }
83
84    /// Clears all entries from the index.
85    pub fn clear(&self) {
86        self.map.write().clear();
87    }
88}
89
90impl<K: Ord + Clone, V: Copy> Default for BTreeIndex<K, V> {
91    fn default() -> Self {
92        Self::new()
93    }
94}
95
96/// A BTree index from i64 keys to NodeIds.
97pub type Int64Index = BTreeIndex<i64, NodeId>;
98
99/// A BTree index from f64 keys to NodeIds.
100/// Note: f64 doesn't implement Ord, so we use a wrapper.
101pub type Float64Index = BTreeIndex<OrderedFloat, NodeId>;
102
103/// A wrapper around f64 that implements Ord.
104#[derive(Debug, Clone, Copy, PartialEq)]
105pub struct OrderedFloat(pub f64);
106
107impl Eq for OrderedFloat {}
108
109impl PartialOrd for OrderedFloat {
110    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
111        Some(self.cmp(other))
112    }
113}
114
115impl Ord for OrderedFloat {
116    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
117        self.0
118            .partial_cmp(&other.0)
119            .unwrap_or(std::cmp::Ordering::Equal)
120    }
121}
122
123impl From<f64> for OrderedFloat {
124    fn from(f: f64) -> Self {
125        Self(f)
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132
133    #[test]
134    fn test_btree_basic() {
135        let index: Int64Index = BTreeIndex::new();
136
137        index.insert(10, NodeId::new(100));
138        index.insert(20, NodeId::new(200));
139        index.insert(30, NodeId::new(300));
140
141        assert_eq!(index.get(&10), Some(NodeId::new(100)));
142        assert_eq!(index.get(&20), Some(NodeId::new(200)));
143        assert_eq!(index.get(&15), None);
144    }
145
146    #[test]
147    fn test_btree_range() {
148        let index: Int64Index = BTreeIndex::new();
149
150        index.insert(10, NodeId::new(100));
151        index.insert(20, NodeId::new(200));
152        index.insert(30, NodeId::new(300));
153        index.insert(40, NodeId::new(400));
154
155        let range = index.range(15..=35);
156        assert_eq!(range.len(), 2);
157        assert!(range.contains(&(20, NodeId::new(200))));
158        assert!(range.contains(&(30, NodeId::new(300))));
159    }
160
161    #[test]
162    fn test_btree_min_max() {
163        let index: Int64Index = BTreeIndex::new();
164
165        assert!(index.min().is_none());
166        assert!(index.max().is_none());
167
168        index.insert(20, NodeId::new(200));
169        index.insert(10, NodeId::new(100));
170        index.insert(30, NodeId::new(300));
171
172        assert_eq!(index.min(), Some((10, NodeId::new(100))));
173        assert_eq!(index.max(), Some((30, NodeId::new(300))));
174    }
175
176    #[test]
177    fn test_btree_remove() {
178        let index: Int64Index = BTreeIndex::new();
179
180        index.insert(10, NodeId::new(100));
181        index.insert(20, NodeId::new(200));
182
183        let removed = index.remove(&10);
184        assert_eq!(removed, Some(NodeId::new(100)));
185        assert!(!index.contains(&10));
186        assert_eq!(index.len(), 1);
187    }
188
189    #[test]
190    fn test_float64_index() {
191        let index: Float64Index = BTreeIndex::new();
192
193        index.insert(1.5.into(), NodeId::new(100));
194        index.insert(2.5.into(), NodeId::new(200));
195        index.insert(3.5.into(), NodeId::new(300));
196
197        let range = index.range(OrderedFloat(1.0)..OrderedFloat(3.0));
198        assert_eq!(range.len(), 2);
199    }
200}
graphos_core/index/btree.rs

graphos_core/index/
btree.rs
