graph_api_simplegraph/index/
range.rs

1use std::borrow::Borrow;
2use std::collections::{BTreeMap, HashSet};
3use std::fmt::Debug;
4use std::hash::Hash;
5use std::ops::RangeBounds;
6
7/// This range index allows range searches for keys.
8#[derive(Debug, Default)]
9pub(crate) struct RangeIndex<K, V>
10where
11    K: Ord + Clone,
12    V: Hash + Eq,
13{
14    map: BTreeMap<K, HashSet<V>>,
15    empty: HashSet<V>,
16}
17
18impl<K, V> RangeIndex<K, V>
19where
20    K: Eq + Hash + Clone + Debug + Ord,
21    V: Eq + Hash + Clone + Copy + Debug,
22{
23    pub(crate) fn insert(&mut self, key: K, value: V) -> bool {
24        self.map.entry(key).or_default().insert(value)
25    }
26
27    pub(crate) fn remove<Q>(&mut self, key: &Q, value: &V) -> bool
28    where
29        K: Borrow<Q> + Ord,
30        Q: ?Sized + Ord,
31    {
32        if let Some(values) = self.map.get_mut(key) {
33            let removed = values.remove(value);
34            if values.is_empty() {
35                self.map.remove(key);
36            }
37            return removed;
38        }
39        false
40    }
41
42    pub(crate) fn get<'a, Q>(&'a self, key: &Q) -> impl Iterator<Item = V> + 'a
43    where
44        K: Borrow<Q>,
45        Q: ?Sized + Hash + Eq + Ord,
46    {
47        self.map.get(key).unwrap_or(&self.empty).iter().copied()
48    }
49
50    pub(crate) fn range<T, R>(&self, range: R) -> impl Iterator<Item = V> + '_
51    where
52        T: ?Sized + Ord,
53        K: Borrow<T> + Ord,
54        R: RangeBounds<T>,
55    {
56        self.map.range(range).flat_map(|(_, v)| v.iter()).copied()
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn test_range_index() {
66        let mut index = RangeIndex::default();
67
68        // Test insertion
69        assert!(index.insert(1, "a"));
70        assert!(index.insert(1, "b"));
71        assert!(index.insert(2, "c"));
72
73        // Test retrieval
74        assert_eq!(index.get(&1).count(), 2);
75        assert_eq!(index.get(&2).count(), 1);
76
77        // Test range query
78        let range_results: Vec<_> = index.range(1..3).collect();
79        assert_eq!(range_results.len(), 3);
80
81        // Test removal
82        assert!(index.remove(&1, &"a"));
83        assert_eq!(index.get(&1).count(), 1);
84
85        // Test complete removal of key when no values remain
86        assert!(index.remove(&1, &"b"));
87        assert_eq!(index.get(&1).count(), 0);
88    }
89
90    #[test]
91    fn test_duplicate_prevention() {
92        let mut index = RangeIndex::default();
93
94        // First insertion should succeed
95        assert!(index.insert(1, "a"));
96
97        // Second insertion of same key-value should fail
98        assert!(!index.insert(1, "a"));
99
100        // Verify only one value exists
101        assert_eq!(index.get(&1).count(), 1);
102
103        // Different value with same key should succeed
104        assert!(index.insert(1, "b"));
105        assert_eq!(index.get(&1).count(), 2);
106    }
107}