intmap/
entry.rs

1// ***************** Entry *********************
2
3use crate::{int::SealedInt, IntKey, IntMap};
4
5/// A view into a single entry in a [`IntMap`], which may either be vacant or occupied.
6///
7/// The entry can be constructed by calling [`IntMap::entry`] with a key. It allows inspection
8/// and in-place manipulation of its value without repeated lookups.
9pub enum Entry<'a, K: IntKey, V: 'a> {
10    /// The entry is occupied.
11    Occupied(OccupiedEntry<'a, K, V>),
12    /// The entry is vacant.
13    Vacant(VacantEntry<'a, K, V>),
14}
15
16impl<'a, K: IntKey, V> Entry<'a, K, V> {
17    #[inline]
18    pub(crate) fn new(key: K, int_map: &'a mut IntMap<K, V>) -> Self {
19        let (cache_ix, vals_ix) = Self::indices(key, int_map);
20
21        match vals_ix {
22            Some(vals_ix) => Entry::Occupied(OccupiedEntry {
23                vals_ix,
24                vals: &mut int_map.cache[cache_ix],
25                count: &mut int_map.count,
26            }),
27            None => Entry::Vacant(VacantEntry {
28                key,
29                cache_ix,
30                int_map,
31            }),
32        }
33    }
34
35    fn indices(key: K, int_map: &IntMap<K, V>) -> (usize, Option<usize>) {
36        if int_map.cache.is_empty() {
37            // Returning 0 is okay because we'll increase the cache and recalculate the index if the
38            // user calls `insert`.
39            return (0, None);
40        }
41
42        let k = key.into_int();
43        let cache_ix = k.calc_index(int_map.mod_mask, K::PRIME);
44
45        let vals = &int_map.cache[cache_ix];
46        let vals_ix = vals.iter().position(|(key, _)| key.into_int() == k);
47
48        (cache_ix, vals_ix)
49    }
50
51    /// Ensures a value is in the entry by inserting the provided value if empty, and returns
52    /// a mutable reference to the value in the entry.
53    pub fn or_insert(self, default: V) -> &'a mut V {
54        match self {
55            Entry::Occupied(entry) => entry.into_mut(),
56            Entry::Vacant(entry) => entry.insert(default),
57        }
58    }
59
60    /// Ensures a value is in the entry by inserting the result of the provided function if empty,
61    /// and returns a mutable reference to the value in the entry.
62    pub fn or_insert_with<F>(self, default: F) -> &'a mut V
63    where
64        F: FnOnce() -> V,
65    {
66        match self {
67            Entry::Occupied(entry) => entry.into_mut(),
68            Entry::Vacant(entry) => entry.insert(default()),
69        }
70    }
71
72    /// Ensures a value is in the entry by inserting, if empty, the result of the provided function.
73    pub fn or_insert_with_key<F>(self, default: F) -> &'a mut V
74    where
75        F: FnOnce(K) -> V,
76    {
77        match self {
78            Entry::Occupied(entry) => entry.into_mut(),
79            Entry::Vacant(entry) => {
80                let d = default(entry.key);
81                entry.insert(d)
82            }
83        }
84    }
85}
86
87impl<'a, K: IntKey, V> Entry<'a, K, V>
88where
89    V: Default,
90{
91    /// Ensures a value is in the entry by inserting the default value if empty,
92    /// and returns a mutable reference to the value in the entry.
93    pub fn or_default(self) -> &'a mut V {
94        match self {
95            Entry::Occupied(entry) => entry.into_mut(),
96            Entry::Vacant(entry) => entry.insert(Default::default()),
97        }
98    }
99}
100
101/// A view into an occupied entry in a [`IntMap`]. It is part of the [`Entry`] enum.
102pub struct OccupiedEntry<'a, K: IntKey, V: 'a> {
103    // Index to vals, guaranteed to be valid
104    vals_ix: usize,
105    // Element of IntMap::cache, guaranteed to be non-empty
106    vals: &'a mut Vec<(K, V)>,
107    // IntMap::count, guaranteed to be non-zero
108    count: &'a mut usize,
109}
110
111impl<'a, K: IntKey, V> OccupiedEntry<'a, K, V> {
112    /// Gets a reference to the value in the entry.
113    pub fn get(&self) -> &V {
114        // Safety: We didn't modify the cache since we calculated the index
115        &self.vals.get(self.vals_ix).unwrap().1
116    }
117
118    /// Gets a mutable reference to the value in the entry.
119    pub fn get_mut(&mut self) -> &mut V {
120        // Safety: We didn't modify the cache since we calculated the index
121        &mut self.vals.get_mut(self.vals_ix).unwrap().1
122    }
123
124    /// Converts the entry into a mutable reference to the value in the entry with a
125    /// lifetime bound to the [`IntMap`] itself.
126    pub fn into_mut(self) -> &'a mut V {
127        // Safety: We didn't modify the cache since we calculated the index
128        &mut self.vals.get_mut(self.vals_ix).unwrap().1
129    }
130
131    /// Sets the value of the entry and returns the old value.
132    pub fn insert(&mut self, value: V) -> V {
133        std::mem::replace(&mut self.vals[self.vals_ix].1, value)
134    }
135
136    /// Removes the value out of the entry and returns it.
137    pub fn remove(self) -> V {
138        // Warning: We modify the cache here, so the index is now invalid
139        *self.count -= 1;
140        let kv = self.vals.swap_remove(self.vals_ix);
141
142        kv.1
143    }
144}
145
146/// A view into a vacant entry in a [`IntMap`]. It is part of the [`Entry`] enum.
147pub struct VacantEntry<'a, K: IntKey, V: 'a> {
148    key: K,
149    cache_ix: usize,
150    int_map: &'a mut IntMap<K, V>,
151}
152
153impl<'a, K: IntKey, V: 'a> VacantEntry<'a, K, V> {
154    pub fn insert(mut self, value: V) -> &'a mut V {
155        if self.int_map.increase_cache_if_needed() {
156            // Recompute cache_ix for the new size.
157            let k = self.key.into_int();
158            self.cache_ix = k.calc_index(self.int_map.mod_mask, K::PRIME);
159        }
160
161        self.int_map.count += 1;
162        let vals = &mut self.int_map.cache[self.cache_ix];
163        vals.push((self.key, value));
164        &mut vals.last_mut().unwrap().1
165    }
166}
intmap/entry.rs

intmap/
entry.rs
