devices/virtio/fs/
multikey.rs

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

devices/virtio/fs/
multikey.rs
