short_lease_map/
lib.rs

1use std::time::{Duration, Instant};
2
3/// A HashMap like collection, but optimized for really short term internship.
4///
5/// It's easiest to think of this like a hotel. When you check in, a room number
6/// is assigned to you. When you leave, that room can now be assigned to someone else.
7#[derive(Clone, Debug)]
8pub struct ShortLeaseMap<T>(Vec<Option<(T, Instant)>>);
9
10impl<T> ShortLeaseMap<T> {
11    /// Creates a new ShortLeaseMap with zero capacity. Capacity will grow as items are added.
12    pub fn new() -> Self {
13        Self::default()
14    }
15
16    /// Creates a new ShortLeaseMap with space reserved for `size` entries.
17    pub fn with_capacity(size: usize) -> Self {
18        Self(Vec::with_capacity(size))
19    }
20
21    /// Adds a value to the map. The value returned can later be used to retrieve it. The returned
22    /// key is not guaranteed to be unique once the value has been removed from this map.
23    pub fn insert(&mut self, t: T) -> usize {
24        let idx = self.0.iter_mut().enumerate().find(|(_, v)| v.is_none());
25        match idx {
26            None => {
27                let idx = self.0.len();
28                self.0.push(Some((t, Instant::now())));
29                idx
30            }
31            Some((i, v)) => {
32                *v = Some((t, Instant::now()));
33                i
34            }
35        }
36    }
37
38    pub fn get(&self, idx: usize) -> Option<&T> {
39        Option::flatten(self.0.get(idx).map(Option::as_ref)).map(|o| &o.0)
40    }
41
42    /// Removes the value with this index. The index may be assigned again after it has been removed.
43    pub fn remove(&mut self, idx: usize) -> Option<T> {
44        self.0.get_mut(idx).and_then(Option::take).map(|o| o.0)
45    }
46
47    /// Evict guests which have overstayed their welcome. If a value has been in the map longer than
48    /// the `max_age` given, it will be dropped. Returns a count of how many items were removed.
49    ///
50    // Clippy will suggest we simplify this code by using Iterator::flatten. It is wrong, the code
51    // is not able to change the iterated value to `None` while using Iterator::flatten.
52    #[allow(clippy::manual_flatten)]
53    pub fn dump_old_values(&mut self, max_age: Duration) -> usize {
54        let mut total_dumped = 0;
55        for e in &mut self.0 {
56            if let Some((_, insert_time)) = e {
57                if insert_time.elapsed() > max_age {
58                    *e = None;
59                    total_dumped += 1;
60                }
61            }
62        }
63        total_dumped
64    }
65
66    /// Iterates immutably over the collection, returning a tuple of a reference to the item and its
67    /// ID value.
68    pub fn iter(&self) -> impl Iterator<Item = (&T, usize)> + DoubleEndedIterator {
69        self.0
70            .iter()
71            .enumerate()
72            .filter_map(|(i, e)| (e.as_ref().map(|o| (&o.0, i))))
73    }
74
75    /// Iterates mutably over the collection, returning a tuple of a mutable reference to the item
76    /// and its ID value.
77    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&mut T, usize)> + DoubleEndedIterator {
78        self.0
79            .iter_mut()
80            .enumerate()
81            .filter_map(|(i, e)| e.as_mut().map(|o| (&mut o.0, i)))
82    }
83}
84
85impl<T> Default for ShortLeaseMap<T> {
86    fn default() -> Self {
87        Self(Vec::default())
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn short_lease_map() {
97        const CAPACITY: usize = 10;
98        let mut map = ShortLeaseMap::with_capacity(CAPACITY);
99        assert_eq!(map.0.capacity(), CAPACITY);
100        for i in 0..CAPACITY + 1 {
101            assert_eq!(map.insert(i), i);
102        }
103        assert_eq!(map.remove(3), Some(3));
104        assert_eq!(map.insert(0), 3);
105        assert_eq!(map.insert(5), CAPACITY + 1);
106        assert_eq!(map.remove(3), Some(0));
107        assert_eq!(map.insert(0), 3);
108    }
109}