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}