graph_api_simplegraph/index/
range.rs1use std::borrow::Borrow;
2use std::collections::{BTreeMap, HashSet};
3use std::fmt::Debug;
4use std::hash::Hash;
5use std::ops::RangeBounds;
6
7#[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 assert!(index.insert(1, "a"));
70 assert!(index.insert(1, "b"));
71 assert!(index.insert(2, "c"));
72
73 assert_eq!(index.get(&1).count(), 2);
75 assert_eq!(index.get(&2).count(), 1);
76
77 let range_results: Vec<_> = index.range(1..3).collect();
79 assert_eq!(range_results.len(), 3);
80
81 assert!(index.remove(&1, &"a"));
83 assert_eq!(index.get(&1).count(), 1);
84
85 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 assert!(index.insert(1, "a"));
96
97 assert!(!index.insert(1, "a"));
99
100 assert_eq!(index.get(&1).count(), 1);
102
103 assert!(index.insert(1, "b"));
105 assert_eq!(index.get(&1).count(), 2);
106 }
107}