architect_api/utils/
bimap.rs

1use std::{borrow::Borrow, collections::BTreeMap};
2
3#[derive(Debug, Clone)]
4pub struct Bimap<U: Clone + Ord, V: Clone + Ord> {
5    map: BTreeMap<U, V>,
6    rev: BTreeMap<V, U>,
7}
8
9impl<U: Clone + Ord, V: Clone + Ord> Bimap<U, V> {
10    pub fn new() -> Self {
11        Self { map: BTreeMap::new(), rev: BTreeMap::new() }
12    }
13
14    pub fn insert(&mut self, u: U, v: V) {
15        self.remove(&u);
16        self.remove_rev(&v);
17        self.map.insert(u.clone(), v.clone());
18        self.rev.insert(v, u);
19    }
20
21    pub fn get(&self, u: &U) -> Option<&V> {
22        self.map.get(u)
23    }
24
25    pub fn get_rev<Q>(&self, v: &Q) -> Option<&U>
26    where
27        Q: ?Sized + Ord,
28        V: Borrow<Q> + Ord,
29    {
30        self.rev.get(v)
31    }
32
33    pub fn get_mut(&mut self, u: &U) -> Option<&mut V> {
34        self.map.get_mut(u)
35    }
36
37    pub fn get_rev_mut(&mut self, v: &V) -> Option<&mut U> {
38        self.rev.get_mut(v)
39    }
40
41    pub fn contains(&self, u: &U) -> bool {
42        self.map.contains_key(u)
43    }
44
45    pub fn contains_rev(&self, v: &V) -> bool {
46        self.rev.contains_key(v)
47    }
48
49    pub fn len(&self) -> usize {
50        self.map.len()
51    }
52
53    pub fn is_empty(&self) -> bool {
54        self.map.is_empty()
55    }
56
57    pub fn remove(&mut self, u: &U) -> Option<V> {
58        let v = self.map.remove(u)?;
59        self.rev.remove(&v);
60        Some(v)
61    }
62
63    pub fn remove_rev(&mut self, v: &V) -> Option<U> {
64        let u = self.rev.remove(v)?;
65        self.map.remove(&u);
66        Some(u)
67    }
68
69    pub fn clear(&mut self) {
70        self.map.clear();
71        self.rev.clear();
72    }
73
74    pub fn retain(&mut self, mut f: impl FnMut(&U, &V) -> bool) {
75        self.map.retain(|u, v| f(u, v));
76        self.rev.retain(|v, u| f(u, v));
77    }
78
79    pub fn iter(&self) -> impl Iterator<Item = (&U, &V)> {
80        self.map.iter()
81    }
82}
83
84impl<U: Clone + Ord, V: Clone + Ord> Default for Bimap<U, V> {
85    fn default() -> Self {
86        Self::new()
87    }
88}
89
90impl<U: Clone + Ord, V: Clone + Ord> FromIterator<(U, V)> for Bimap<U, V> {
91    fn from_iter<I: IntoIterator<Item = (U, V)>>(iter: I) -> Self {
92        let mut bimap = Bimap::new();
93        for (u, v) in iter {
94            bimap.insert(u, v);
95        }
96        bimap
97    }
98}
99
100impl<U: Clone + Ord, V: Clone + Ord> IntoIterator for Bimap<U, V> {
101    type Item = (U, V);
102    type IntoIter = std::collections::btree_map::IntoIter<U, V>;
103
104    fn into_iter(self) -> Self::IntoIter {
105        self.map.into_iter()
106    }
107}