shared/
bijective_map.rs

1use core::fmt::Debug;
2
3use alloc::collections::btree_map::BTreeMap;
4
5pub struct BijectiveBTreeMap<L, R> {
6    to_left: BTreeMap<R, L>,
7    to_right: BTreeMap<L, R>,
8}
9
10impl<L: Ord + Copy, R: Ord + Copy> BijectiveBTreeMap<L, R> {
11    pub const fn new() -> Self {
12        Self {
13            to_left: BTreeMap::new(),
14            to_right: BTreeMap::new(),
15        }
16    }
17
18    pub fn insert(&mut self, key: L, value: R) {
19        self.to_left.insert(value, key);
20        self.to_right.insert(key, value);
21    }
22
23    pub fn get_by_left(&self, key: &L) -> Option<&R> {
24        self.to_right.get(key)
25    }
26
27    pub fn get_by_right(&self, value: &R) -> Option<&L> {
28        self.to_left.get(value)
29    }
30
31    pub fn remove_by_key(&mut self, key: &L) -> Option<R> {
32        if let Some(value) = self.to_right.remove(key) {
33            self.to_left.remove(&value);
34            Some(value)
35        } else {
36            None
37        }
38    }
39
40    pub fn remove_by_value(&mut self, value: &R) -> Option<L> {
41        if let Some(key) = self.to_left.remove(value) {
42            self.to_right.remove(&key);
43            Some(key)
44        } else {
45            None
46        }
47    }
48
49    pub fn get_left_keys(&self) -> impl Iterator<Item = &L> {
50        self.to_right.keys()
51    }
52
53    pub fn get_right_keys(&self) -> impl Iterator<Item = &R> {
54        self.to_left.keys()
55    }
56}
57
58impl<L: Ord + Copy, R: Ord + Copy> Default for BijectiveBTreeMap<L, R> {
59    fn default() -> Self {
60        Self::new()
61    }
62}
63
64impl<L: Ord + Copy + Debug, R: Ord + Copy + Debug> Debug for BijectiveBTreeMap<L, R> {
65    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
66        f.debug_struct("BijectiveBTreeMap")
67            .field("to_key", &self.to_left)
68            .field("to_value", &self.to_right)
69            .finish()
70    }
71}
72
73impl<L: Ord + Copy, R: Ord + Copy> Clone for BijectiveBTreeMap<L, R> {
74    fn clone(&self) -> Self {
75        Self {
76            to_left: self.to_left.clone(),
77            to_right: self.to_right.clone(),
78        }
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use alloc::{format, vec, vec::Vec};
85
86    use super::*;
87
88    #[test]
89    fn test_new() {
90        let map: BijectiveBTreeMap<i32, i32> = BijectiveBTreeMap::new();
91        assert!(map.get_by_left(&1).is_none());
92        assert!(map.get_by_right(&1).is_none());
93    }
94
95    #[test]
96    fn test_insert_and_get() {
97        let mut map = BijectiveBTreeMap::new();
98        map.insert(1, 10);
99        map.insert(2, 20);
100
101        assert_eq!(map.get_by_left(&1), Some(&10));
102        assert_eq!(map.get_by_left(&2), Some(&20));
103        assert_eq!(map.get_by_right(&10), Some(&1));
104        assert_eq!(map.get_by_right(&20), Some(&2));
105    }
106
107    #[test]
108    fn test_insert_overwrite() {
109        let mut map = BijectiveBTreeMap::new();
110        map.insert(1, 10);
111        map.insert(1, 20);
112
113        assert_eq!(map.get_by_left(&1), Some(&20));
114        assert!(map.get_by_right(&10).is_some());
115    }
116
117    #[test]
118    fn test_remove_by_key() {
119        let mut map = BijectiveBTreeMap::new();
120        map.insert(1, 10);
121        map.insert(2, 20);
122
123        assert_eq!(map.remove_by_key(&1), Some(10));
124        assert!(map.get_by_left(&1).is_none());
125        assert!(map.get_by_right(&10).is_none());
126        assert_eq!(map.get_by_left(&2), Some(&20));
127    }
128
129    #[test]
130    fn test_remove_by_value() {
131        let mut map = BijectiveBTreeMap::new();
132        map.insert(1, 10);
133        map.insert(2, 20);
134
135        assert_eq!(map.remove_by_value(&10), Some(1));
136        assert!(map.get_by_left(&1).is_none());
137        assert!(map.get_by_right(&10).is_none());
138        assert_eq!(map.get_by_left(&2), Some(&20));
139    }
140
141    #[test]
142    fn test_remove_nonexistent() {
143        let mut map = BijectiveBTreeMap::new();
144        map.insert(1, 10);
145
146        assert_eq!(map.remove_by_key(&2), None);
147        assert_eq!(map.remove_by_value(&20), None);
148    }
149
150    #[test]
151    fn test_iterators() {
152        let mut map = BijectiveBTreeMap::new();
153        map.insert(1, 10);
154        map.insert(2, 20);
155        map.insert(3, 30);
156
157        let left_keys: Vec<_> = map.get_left_keys().copied().collect();
158        assert_eq!(left_keys, vec![1, 2, 3]);
159
160        let right_keys: Vec<_> = map.get_right_keys().copied().collect();
161        assert_eq!(right_keys, vec![10, 20, 30]);
162    }
163
164    #[test]
165    fn test_clone() {
166        let mut map = BijectiveBTreeMap::new();
167        map.insert(1, 10);
168        map.insert(2, 20);
169
170        let cloned = map.clone();
171        assert_eq!(cloned.get_by_left(&1), Some(&10));
172        assert_eq!(cloned.get_by_right(&20), Some(&2));
173    }
174
175    #[test]
176    fn test_default() {
177        let map: BijectiveBTreeMap<i32, i32> = BijectiveBTreeMap::default();
178        assert!(map.get_by_left(&1).is_none());
179    }
180
181    #[test]
182    fn test_debug() {
183        let mut map = BijectiveBTreeMap::new();
184        map.insert(1, 10);
185        let debug_str = format!("{:?}", map);
186        assert!(debug_str.contains("BijectiveBTreeMap"));
187    }
188}