take_some/
btree_map.rs

1use std::collections::BTreeMap;
2use TakeSome;
3
4impl<K, V> TakeSome for BTreeMap<K, V>
5where
6    K: Ord,
7{
8    type Item = (K, V);
9
10    /// Takes the "greatest" elements from the map.
11    ///
12    /// `unsafe` under the cover!
13    fn take_some(&mut self) -> Option<(K, V)> {
14        // If the key corresponds to the "first" element (the one we store separately), we have
15        // to take a "random" element from the map (`self.rest`) and put is as the "first".
16        let temp_max_key: *const K = match self.keys().next_back() {
17            Some(x) => x,
18            None => return None,
19        };
20        // We need that unsafe magic because we are going to remove an element a reference to
21        // which we kinda hold.
22        let temp_max_key: &K = unsafe { &*temp_max_key };
23        let new_map = self.split_off(temp_max_key);
24        // After a split we have a `new_map` which contains "everything after a given key,
25        // including the key", and because it's a maximum key, the `new_map` contains of only
26        // that key.
27        Some(
28            new_map
29                .into_iter()
30                .next()
31                .expect("Well, it should contain a key"),
32        )
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39    use std::iter::once;
40
41    #[test]
42    fn check_btree_map() {
43        let mut empty: BTreeMap<String, String> = BTreeMap::new();
44        assert_eq!(None, empty.take_some());
45
46        let mut one_element: BTreeMap<_, _> =
47            once(("check".to_string(), "checked".to_string())).collect();
48        assert_eq!(
49            Some(("check".into(), "checked".into())),
50            one_element.take_some()
51        );
52        assert!(one_element.is_empty());
53
54        let mut map: BTreeMap<_, _> = vec![
55            ("check".to_string(), "lol".to_string()),
56            ("wut".into(), "ahaha".into()),
57            ("123123".into(), "....".into()),
58        ].into_iter()
59            .collect();
60        assert_eq!(Some(("wut".into(), "ahaha".into())), map.take_some());
61        assert_eq!(2, map.len());
62    }
63}