libtw2_net/collections/
peer_map.rs

1use crate::net::PeerId;
2use linear_map::LinearMap;
3use std::fmt;
4use std::iter::FromIterator;
5use std::ops;
6
7#[derive(Clone)]
8pub struct PeerMap<T> {
9    // TODO: Different data structure, HashMap?
10    map: LinearMap<PeerId, T>,
11}
12
13impl<T> Default for PeerMap<T> {
14    fn default() -> PeerMap<T> {
15        PeerMap::new()
16    }
17}
18
19impl<T: fmt::Debug> fmt::Debug for PeerMap<T> {
20    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
21        f.debug_map().entries(self.iter()).finish()
22    }
23}
24
25impl<T> PeerMap<T> {
26    pub fn new() -> PeerMap<T> {
27        PeerMap {
28            map: LinearMap::new(),
29        }
30    }
31    pub fn with_capacity(cap: usize) -> PeerMap<T> {
32        PeerMap {
33            map: LinearMap::with_capacity(cap),
34        }
35    }
36    pub fn clear(&mut self) {
37        self.map.clear()
38    }
39    pub fn len(&self) -> usize {
40        self.map.len()
41    }
42    pub fn is_empty(&self) -> bool {
43        self.map.is_empty()
44    }
45    pub fn iter(&self) -> Iter<T> {
46        Iter(self.map.iter())
47    }
48    pub fn iter_mut(&mut self) -> IterMut<T> {
49        IterMut(self.map.iter_mut())
50    }
51    pub fn keys(&self) -> Keys<T> {
52        Keys(self.map.keys())
53    }
54    pub fn values(&self) -> Values<T> {
55        Values(self.map.values())
56    }
57    pub fn drain(&mut self) -> Drain<T> {
58        Drain(self.map.drain())
59    }
60    pub fn insert(&mut self, pid: PeerId, value: T) -> Option<T> {
61        self.map.insert(pid, value)
62    }
63    pub fn remove(&mut self, pid: PeerId) {
64        self.map
65            .remove(&pid)
66            .unwrap_or_else(|| panic!("invalid pid"));
67    }
68    pub fn get(&self, pid: PeerId) -> Option<&T> {
69        self.map.get(&pid)
70    }
71    pub fn get_mut(&mut self, pid: PeerId) -> Option<&mut T> {
72        self.map.get_mut(&pid)
73    }
74    pub fn entry(&mut self, pid: PeerId) -> Entry<T> {
75        match self.map.entry(pid) {
76            linear_map::Entry::Occupied(o) => Entry::Occupied(OccupiedEntry(o)),
77            linear_map::Entry::Vacant(v) => Entry::Vacant(VacantEntry(v)),
78        }
79    }
80    pub fn contains_key(&mut self, pid: PeerId) -> bool {
81        self.map.contains_key(&pid)
82    }
83}
84
85impl<T> ops::Index<PeerId> for PeerMap<T> {
86    type Output = T;
87    fn index(&self, pid: PeerId) -> &T {
88        self.get(pid).unwrap_or_else(|| panic!("invalid pid"))
89    }
90}
91
92impl<T> ops::IndexMut<PeerId> for PeerMap<T> {
93    fn index_mut(&mut self, pid: PeerId) -> &mut T {
94        self.get_mut(pid).unwrap_or_else(|| panic!("invalid pid"))
95    }
96}
97
98impl<T> FromIterator<(PeerId, T)> for PeerMap<T> {
99    fn from_iter<I>(iter: I) -> PeerMap<T>
100    where
101        I: IntoIterator<Item = (PeerId, T)>,
102    {
103        PeerMap {
104            map: FromIterator::from_iter(iter),
105        }
106    }
107}
108
109impl<T> Extend<(PeerId, T)> for PeerMap<T> {
110    fn extend<I>(&mut self, iter: I)
111    where
112        I: IntoIterator<Item = (PeerId, T)>,
113    {
114        self.map.extend(iter)
115    }
116}
117
118impl<'a, T: 'a> IntoIterator for &'a PeerMap<T> {
119    type Item = (PeerId, &'a T);
120    type IntoIter = Iter<'a, T>;
121    fn into_iter(self) -> Iter<'a, T> {
122        self.iter()
123    }
124}
125
126impl<'a, T: 'a> IntoIterator for &'a mut PeerMap<T> {
127    type Item = (PeerId, &'a mut T);
128    type IntoIter = IterMut<'a, T>;
129    fn into_iter(self) -> IterMut<'a, T> {
130        self.iter_mut()
131    }
132}
133
134pub struct Iter<'a, T: 'a>(linear_map::Iter<'a, PeerId, T>);
135pub struct IterMut<'a, T: 'a>(linear_map::IterMut<'a, PeerId, T>);
136pub struct Drain<'a, T: 'a>(linear_map::Drain<'a, PeerId, T>);
137pub struct Keys<'a, T: 'a>(linear_map::Keys<'a, PeerId, T>);
138pub struct Values<'a, T: 'a>(linear_map::Values<'a, PeerId, T>);
139
140impl<'a, T> Iterator for Iter<'a, T> {
141    type Item = (PeerId, &'a T);
142    fn next(&mut self) -> Option<(PeerId, &'a T)> {
143        self.0.next().map(|(&pid, e)| (pid, e))
144    }
145    fn size_hint(&self) -> (usize, Option<usize>) {
146        self.0.size_hint()
147    }
148}
149
150impl<'a, T> Iterator for IterMut<'a, T> {
151    type Item = (PeerId, &'a mut T);
152    fn next(&mut self) -> Option<(PeerId, &'a mut T)> {
153        self.0.next().map(|(&pid, e)| (pid, e))
154    }
155    fn size_hint(&self) -> (usize, Option<usize>) {
156        self.0.size_hint()
157    }
158}
159
160impl<'a, T> Iterator for Drain<'a, T> {
161    type Item = (PeerId, T);
162    fn next(&mut self) -> Option<(PeerId, T)> {
163        self.0.next()
164    }
165    fn size_hint(&self) -> (usize, Option<usize>) {
166        self.0.size_hint()
167    }
168}
169
170impl<'a, T: 'a> Iterator for Keys<'a, T> {
171    type Item = PeerId;
172    fn next(&mut self) -> Option<PeerId> {
173        self.0.next().cloned()
174    }
175    fn size_hint(&self) -> (usize, Option<usize>) {
176        self.0.size_hint()
177    }
178}
179
180impl<'a, T: 'a> Iterator for Values<'a, T> {
181    type Item = &'a T;
182    fn next(&mut self) -> Option<&'a T> {
183        self.0.next()
184    }
185    fn size_hint(&self) -> (usize, Option<usize>) {
186        self.0.size_hint()
187    }
188}
189
190pub struct OccupiedEntry<'a, T: 'a>(linear_map::OccupiedEntry<'a, PeerId, T>);
191pub struct VacantEntry<'a, T: 'a>(linear_map::VacantEntry<'a, PeerId, T>);
192
193pub enum Entry<'a, T: 'a> {
194    Occupied(OccupiedEntry<'a, T>),
195    Vacant(VacantEntry<'a, T>),
196}
197
198impl<'a, T: 'a> Entry<'a, T> {
199    pub fn assert_occupied(self) -> OccupiedEntry<'a, T> {
200        match self {
201            Entry::Occupied(o) => o,
202            Entry::Vacant(_) => panic!("invalid pid"),
203        }
204    }
205}
206
207impl<'a, T: 'a> OccupiedEntry<'a, T> {
208    pub fn get(&self) -> &T {
209        self.0.get()
210    }
211    pub fn get_mut(&mut self) -> &mut T {
212        self.0.get_mut()
213    }
214    pub fn remove(self) -> T {
215        self.0.remove()
216    }
217}
218
219impl<'a, T: 'a> VacantEntry<'a, T> {
220    pub fn insert(self, value: T) -> &'a mut T {
221        self.0.insert(value)
222    }
223}