1 // Copyright 2019 The Chromium OS Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 use std::borrow::Borrow; 6 use std::collections::BTreeMap; 7 8 /// A BTreeMap that supports 2 types of keys per value. All the usual restrictions and warnings for 9 /// `std::collections::BTreeMap` also apply to this struct. Additionally, there is a 1:1 10 /// relationship between the 2 key types. In other words, for each `K1` in the map, there is exactly 11 /// one `K2` in the map and vice versa. 12 pub struct MultikeyBTreeMap<K1, K2, V> { 13 // We need to keep a copy of the second key in the main map so that we can remove entries using 14 // just the main key. Otherwise we would require the caller to provide both keys when calling 15 // `remove`. 16 main: BTreeMap<K1, (K2, V)>, 17 alt: BTreeMap<K2, K1>, 18 } 19 20 impl<K1, K2, V> MultikeyBTreeMap<K1, K2, V> 21 where 22 K1: Clone + Ord, 23 K2: Clone + Ord, 24 { 25 /// Create a new empty MultikeyBTreeMap. new() -> Self26 pub fn new() -> Self { 27 MultikeyBTreeMap { 28 main: BTreeMap::default(), 29 alt: BTreeMap::default(), 30 } 31 } 32 33 /// Returns a reference to the value corresponding to the key. 34 /// 35 /// The key may be any borrowed form of `K1``, but the ordering on the borrowed form must match 36 /// the ordering on `K1`. get<Q>(&self, key: &Q) -> Option<&V> where K1: Borrow<Q>, Q: Ord + ?Sized,37 pub fn get<Q>(&self, key: &Q) -> Option<&V> 38 where 39 K1: Borrow<Q>, 40 Q: Ord + ?Sized, 41 { 42 self.main.get(key).map(|(_, v)| v) 43 } 44 45 /// Returns a reference to the value corresponding to the alternate key. 46 /// 47 /// The key may be any borrowed form of the `K2``, but the ordering on the borrowed form must 48 /// match the ordering on `K2`. 49 /// 50 /// Note that this method performs 2 lookups: one to get the main key and another to get the 51 /// value associated with that key. For best performance callers should prefer the `get` method 52 /// over this method whenever possible as `get` only needs to perform one lookup. get_alt<Q2>(&self, key: &Q2) -> Option<&V> where K2: Borrow<Q2>, Q2: Ord + ?Sized,53 pub fn get_alt<Q2>(&self, key: &Q2) -> Option<&V> 54 where 55 K2: Borrow<Q2>, 56 Q2: Ord + ?Sized, 57 { 58 if let Some(k) = self.alt.get(key) { 59 self.get(k) 60 } else { 61 None 62 } 63 } 64 65 /// Inserts a new entry into the map with the given keys and value. 66 /// 67 /// Returns `None` if the map did not have an entry with `k1` or `k2` present. If exactly one 68 /// key was present, then the value associated with that key is updated, the other key is 69 /// removed, and the old value is returned. If **both** keys were present then the value 70 /// associated with the main key is updated, the value associated with the alternate key is 71 /// removed, and the old value associated with the main key is returned. insert(&mut self, k1: K1, k2: K2, v: V) -> Option<V>72 pub fn insert(&mut self, k1: K1, k2: K2, v: V) -> Option<V> { 73 let oldval = if let Some(oldkey) = self.alt.insert(k2.clone(), k1.clone()) { 74 self.main.remove(&oldkey) 75 } else { 76 None 77 }; 78 self.main 79 .insert(k1, (k2.clone(), v)) 80 .or(oldval) 81 .map(|(oldk2, v)| { 82 if oldk2 != k2 { 83 self.alt.remove(&oldk2); 84 } 85 v 86 }) 87 } 88 89 /// Remove a key from the map, returning the value associated with that key if it was previously 90 /// in the map. 91 /// 92 /// The key may be any borrowed form of `K1``, but the ordering on the borrowed form must match 93 /// the ordering on `K1`. remove<Q>(&mut self, key: &Q) -> Option<V> where K1: Borrow<Q>, Q: Ord + ?Sized,94 pub fn remove<Q>(&mut self, key: &Q) -> Option<V> 95 where 96 K1: Borrow<Q>, 97 Q: Ord + ?Sized, 98 { 99 self.main.remove(key).map(|(k2, v)| { 100 self.alt.remove(&k2); 101 v 102 }) 103 } 104 105 /// Clears the map, removing all values. clear(&mut self)106 pub fn clear(&mut self) { 107 self.alt.clear(); 108 self.main.clear() 109 } 110 } 111 112 #[cfg(test)] 113 mod test { 114 use super::*; 115 116 #[test] get()117 fn get() { 118 let mut m = MultikeyBTreeMap::<u64, i64, u32>::new(); 119 120 let k1 = 0xc6c8_f5e0_b13e_ed40; 121 let k2 = 0x1a04_ce4b_8329_14fe; 122 let val = 0xf4e3_c360; 123 assert!(m.insert(k1, k2, val).is_none()); 124 125 assert_eq!(*m.get(&k1).expect("failed to look up main key"), val); 126 assert_eq!(*m.get_alt(&k2).expect("failed to look up alt key"), val); 127 } 128 129 #[test] update_main_key()130 fn update_main_key() { 131 let mut m = MultikeyBTreeMap::<u64, i64, u32>::new(); 132 133 let k1 = 0xc6c8_f5e0_b13e_ed40; 134 let k2 = 0x1a04_ce4b_8329_14fe; 135 let val = 0xf4e3_c360; 136 assert!(m.insert(k1, k2, val).is_none()); 137 138 let new_k1 = 0x3add_f8f8_c7c5_df5e; 139 let val2 = 0x7389_f8a7; 140 assert_eq!( 141 m.insert(new_k1, k2, val2) 142 .expect("failed to update main key"), 143 val 144 ); 145 146 assert!(m.get(&k1).is_none()); 147 assert_eq!(*m.get(&new_k1).expect("failed to look up main key"), val2); 148 assert_eq!(*m.get_alt(&k2).expect("failed to look up alt key"), val2); 149 } 150 151 #[test] update_alt_key()152 fn update_alt_key() { 153 let mut m = MultikeyBTreeMap::<u64, i64, u32>::new(); 154 155 let k1 = 0xc6c8_f5e0_b13e_ed40; 156 let k2 = 0x1a04_ce4b_8329_14fe; 157 let val = 0xf4e3_c360; 158 assert!(m.insert(k1, k2, val).is_none()); 159 160 let new_k2 = 0x6825_a60b_61ac_b333; 161 let val2 = 0xbb14_8f2c; 162 assert_eq!( 163 m.insert(k1, new_k2, val2) 164 .expect("failed to update alt key"), 165 val 166 ); 167 168 assert!(m.get_alt(&k2).is_none()); 169 assert_eq!(*m.get(&k1).expect("failed to look up main key"), val2); 170 assert_eq!( 171 *m.get_alt(&new_k2).expect("failed to look up alt key"), 172 val2 173 ); 174 } 175 176 #[test] update_value()177 fn update_value() { 178 let mut m = MultikeyBTreeMap::<u64, i64, u32>::new(); 179 180 let k1 = 0xc6c8_f5e0_b13e_ed40; 181 let k2 = 0x1a04_ce4b_8329_14fe; 182 let val = 0xf4e3_c360; 183 assert!(m.insert(k1, k2, val).is_none()); 184 185 let val2 = 0xe42d_79ba; 186 assert_eq!( 187 m.insert(k1, k2, val2).expect("failed to update alt key"), 188 val 189 ); 190 191 assert_eq!(*m.get(&k1).expect("failed to look up main key"), val2); 192 assert_eq!(*m.get_alt(&k2).expect("failed to look up alt key"), val2); 193 } 194 195 #[test] update_both_keys_main()196 fn update_both_keys_main() { 197 let mut m = MultikeyBTreeMap::<u64, i64, u32>::new(); 198 199 let k1 = 0xc6c8_f5e0_b13e_ed40; 200 let k2 = 0x1a04_ce4b_8329_14fe; 201 let val = 0xf4e3_c360; 202 assert!(m.insert(k1, k2, val).is_none()); 203 204 let new_k1 = 0xc980_587a_24b3_ae30; 205 let new_k2 = 0x2773_c5ee_8239_45a2; 206 let val2 = 0x31f4_33f9; 207 assert!(m.insert(new_k1, new_k2, val2).is_none()); 208 209 let val3 = 0x8da1_9cf7; 210 assert_eq!( 211 m.insert(k1, new_k2, val3) 212 .expect("failed to update main key"), 213 val 214 ); 215 216 // Both new_k1 and k2 should now be gone from the map. 217 assert!(m.get(&new_k1).is_none()); 218 assert!(m.get_alt(&k2).is_none()); 219 220 assert_eq!(*m.get(&k1).expect("failed to look up main key"), val3); 221 assert_eq!( 222 *m.get_alt(&new_k2).expect("failed to look up alt key"), 223 val3 224 ); 225 } 226 227 #[test] update_both_keys_alt()228 fn update_both_keys_alt() { 229 let mut m = MultikeyBTreeMap::<u64, i64, u32>::new(); 230 231 let k1 = 0xc6c8_f5e0_b13e_ed40; 232 let k2 = 0x1a04_ce4b_8329_14fe; 233 let val = 0xf4e3_c360; 234 assert!(m.insert(k1, k2, val).is_none()); 235 236 let new_k1 = 0xc980_587a_24b3_ae30; 237 let new_k2 = 0x2773_c5ee_8239_45a2; 238 let val2 = 0x31f4_33f9; 239 assert!(m.insert(new_k1, new_k2, val2).is_none()); 240 241 let val3 = 0x8da1_9cf7; 242 assert_eq!( 243 m.insert(new_k1, k2, val3) 244 .expect("failed to update main key"), 245 val2 246 ); 247 248 // Both k1 and new_k2 should now be gone from the map. 249 assert!(m.get(&k1).is_none()); 250 assert!(m.get_alt(&new_k2).is_none()); 251 252 assert_eq!(*m.get(&new_k1).expect("failed to look up main key"), val3); 253 assert_eq!(*m.get_alt(&k2).expect("failed to look up alt key"), val3); 254 } 255 256 #[test] remove()257 fn remove() { 258 let mut m = MultikeyBTreeMap::<u64, i64, u32>::new(); 259 260 let k1 = 0xc6c8_f5e0_b13e_ed40; 261 let k2 = 0x1a04_ce4b_8329_14fe; 262 let val = 0xf4e3_c360; 263 assert!(m.insert(k1, k2, val).is_none()); 264 265 assert_eq!(m.remove(&k1).expect("failed to remove entry"), val); 266 assert!(m.get(&k1).is_none()); 267 assert!(m.get_alt(&k2).is_none()); 268 } 269 } 270