take_some/
btree_set.rs

1use std::collections::BTreeSet;
2use TakeSome;
3
4impl<K> TakeSome for BTreeSet<K>
5where
6    K: Ord,
7{
8    type Item = K;
9
10    /// Takes the "greatest" elements from the set.
11    ///
12    /// `unsafe` under the cover!
13    fn take_some(&mut self) -> Option<K> {
14        // If the key corresponds to the "first" element (the one we store separately), we have
15        // to take a "random" element from the set (`self.rest`) and put is as the "first".
16        let temp_key: *const K = match self.iter().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_key: &K = unsafe { &*temp_key };
23        Some(
24            self.take(temp_key)
25                .expect("We've just pulled a key from the set"),
26        )
27    }
28}
29
30#[cfg(test)]
31mod tests {
32    use super::*;
33    use std::iter::once;
34
35    #[test]
36    fn check_btree_set() {
37        let mut empty: BTreeSet<String> = BTreeSet::new();
38        assert_eq!(None, empty.take_some());
39
40        let mut one_element: BTreeSet<_> = once("check".to_string()).collect();
41        assert_eq!(Some("check".into()), one_element.take_some());
42        assert!(one_element.is_empty());
43
44        let mut set: BTreeSet<_> = vec!["check".to_string(), "lol".into(), "wut".into()]
45            .into_iter()
46            .collect();
47        assert_eq!(Some("wut".into()), set.take_some());
48        assert_eq!(2, set.len());
49    }
50}