defuse_map_utils/
cleanup.rs

1// This is due to a bug in clippy that happens with shared security analysis
2// Clippy doesn't seem to play well with autoimpl
3#![allow(clippy::type_repetition_in_bounds)]
4
5use core::{
6    fmt::Debug,
7    marker::PhantomData,
8    ops::{Deref, DerefMut},
9};
10
11use impl_tools::autoimpl;
12
13use crate::{Entry, Map, OccupiedEntry, VacantEntry};
14
15/// A mapping where non-existing keys considered to have [`Default`] values
16pub trait DefaultMap: Map<V: Default + Eq> {
17    /// Get an entry at given key or [`Default`] value if the key doesn't exist.
18    ///
19    /// The returned entry will automatically be removed from the map if it becomes
20    /// equal to [`Default`] after modifications.
21    ///
22    /// # Example
23    ///
24    /// ```rust
25    /// # use std::collections::HashMap;
26    /// # use defuse_map_utils::cleanup::DefaultMap;
27    /// let mut m = HashMap::new();
28    /// *m.entry_or_default("a") += 1;
29    /// assert_eq!(m.get("a"), Some(&1));
30    /// *m.entry_or_default("a") -= 1;
31    /// assert_eq!(m.get("a"), None);
32    /// ```
33    #[inline]
34    fn entry_or_default(
35        &mut self,
36        k: Self::K,
37    ) -> DefaultEntry<'_, Self::VacantEntry<'_>, Self::OccupiedEntry<'_>> {
38        match self.entry(k) {
39            Entry::Vacant(entry) => DefaultVacantEntry::new(entry).into(),
40            Entry::Occupied(entry) => DefaultOccupiedEntry::new(entry).into(),
41        }
42    }
43}
44impl<T> DefaultMap for T where T: Map<V: Default + Eq> {}
45
46#[autoimpl(Debug where V: Debug, V::V: Debug, O: Debug)]
47pub enum DefaultEntry<'a, V, O>
48where
49    V: VacantEntry<'a, V: Default + Eq>,
50    O: OccupiedEntry<'a, K = V::K, V = V::V>,
51{
52    Vacant(DefaultVacantEntry<'a, V>),
53    Occupied(DefaultOccupiedEntry<'a, O>),
54}
55
56impl<'a, V, O> DefaultEntry<'a, V, O>
57where
58    V: VacantEntry<'a, V: Default + Eq>,
59    O: OccupiedEntry<'a, K = V::K, V = V::V>,
60{
61    /// Get the key associated with the entry.
62    ///
63    /// # Example
64    ///
65    /// ```rust
66    /// # use std::collections::HashMap;
67    /// # use defuse_map_utils::cleanup::DefaultMap;
68    /// let mut m: HashMap<_, ()> = HashMap::new();
69    /// assert_eq!(*m.entry_or_default("a").key(), "a");
70    /// ```
71    #[inline]
72    pub fn key(&self) -> &V::K {
73        match self {
74            Self::Vacant(entry) => entry.key(),
75            Self::Occupied(entry) => entry.key(),
76        }
77    }
78
79    /// Remove the entry from the map, regardless of its value.
80    ///
81    /// # Example
82    ///
83    /// ```rust
84    /// # use std::collections::HashMap;
85    /// # use defuse_map_utils::cleanup::DefaultMap;
86    /// let mut m: HashMap<_, i32> = HashMap::new();
87    /// let mut entry = m.entry_or_default("a");
88    /// *entry += 1;
89    /// assert_eq!(entry.remove(), 1);
90    /// assert_eq!(m.get("a"), None);
91    /// ```
92    #[inline]
93    pub fn remove(self) -> V::V {
94        match self {
95            Self::Vacant(entry) => entry.remove(),
96            Self::Occupied(entry) => entry.remove(),
97        }
98    }
99}
100
101impl<'a, V, O> Deref for DefaultEntry<'a, V, O>
102where
103    V: VacantEntry<'a, V: Default + Eq>,
104    O: OccupiedEntry<'a, K = V::K, V = V::V>,
105{
106    type Target = V::V;
107
108    #[inline]
109    fn deref(&self) -> &Self::Target {
110        match self {
111            Self::Vacant(entry) => entry,
112            Self::Occupied(entry) => entry,
113        }
114    }
115}
116
117impl<'a, V, O> DerefMut for DefaultEntry<'a, V, O>
118where
119    V: VacantEntry<'a, V: Default + Eq>,
120    O: OccupiedEntry<'a, K = V::K, V = V::V>,
121{
122    #[inline]
123    fn deref_mut(&mut self) -> &mut Self::Target {
124        match self {
125            Self::Vacant(entry) => entry,
126            Self::Occupied(entry) => entry,
127        }
128    }
129}
130
131impl<'a, V, O> From<DefaultVacantEntry<'a, V>> for DefaultEntry<'a, V, O>
132where
133    V: VacantEntry<'a, V: Default + Eq>,
134    O: OccupiedEntry<'a, K = V::K, V = V::V>,
135{
136    #[inline]
137    fn from(entry: DefaultVacantEntry<'a, V>) -> Self {
138        Self::Vacant(entry)
139    }
140}
141
142impl<'a, V, O> From<DefaultOccupiedEntry<'a, O>> for DefaultEntry<'a, V, O>
143where
144    V: VacantEntry<'a, V: Default + Eq>,
145    O: OccupiedEntry<'a, K = V::K, V = V::V>,
146{
147    #[inline]
148    fn from(entry: DefaultOccupiedEntry<'a, O>) -> Self {
149        Self::Occupied(entry)
150    }
151}
152
153#[derive(Debug)]
154pub struct DefaultVacantEntry<'a, E: 'a>(Option<(E::V, E)>)
155where
156    E: VacantEntry<'a, V: Default + Eq>;
157
158impl<'a, E: 'a> DefaultVacantEntry<'a, E>
159where
160    E: VacantEntry<'a, V: Default + Eq>,
161{
162    #[inline]
163    pub fn new(entry: E) -> Self {
164        Self(Some((Default::default(), entry)))
165    }
166
167    #[inline]
168    pub fn key(&self) -> &E::K {
169        self.0.as_ref().unwrap_or_else(|| unreachable!()).1.key()
170    }
171
172    #[inline]
173    pub fn remove(mut self) -> E::V {
174        self.0.take().unwrap_or_else(|| unreachable!()).0
175    }
176}
177
178impl<'a, E: 'a> Deref for DefaultVacantEntry<'a, E>
179where
180    E: VacantEntry<'a, V: Default + Eq>,
181{
182    type Target = E::V;
183
184    #[inline]
185    fn deref(&self) -> &Self::Target {
186        &self.0.as_ref().unwrap_or_else(|| unreachable!()).0
187    }
188}
189
190impl<'a, E: 'a> DerefMut for DefaultVacantEntry<'a, E>
191where
192    E: VacantEntry<'a, V: Default + Eq>,
193{
194    #[inline]
195    fn deref_mut(&mut self) -> &mut Self::Target {
196        &mut self.0.as_mut().unwrap_or_else(|| unreachable!()).0
197    }
198}
199
200impl<'a, E: 'a> Drop for DefaultVacantEntry<'a, E>
201where
202    E: VacantEntry<'a, V: Default + Eq>,
203{
204    #[inline]
205    fn drop(&mut self) {
206        let Some((v, entry)) = self.0.take() else {
207            return;
208        };
209        if v != Default::default() {
210            entry.insert(v);
211        }
212    }
213}
214
215#[derive(Debug)]
216pub struct DefaultOccupiedEntry<'a, E>(Option<E>, PhantomData<&'a ()>)
217where
218    E: OccupiedEntry<'a, V: Default + Eq>;
219
220impl<'a, E> DefaultOccupiedEntry<'a, E>
221where
222    E: OccupiedEntry<'a, V: Default + Eq>,
223{
224    #[inline]
225    pub const fn new(entry: E) -> Self {
226        Self(Some(entry), PhantomData)
227    }
228
229    #[inline]
230    pub fn key(&self) -> &E::K {
231        self.0.as_ref().unwrap_or_else(|| unreachable!()).key()
232    }
233
234    #[inline]
235    pub fn remove(mut self) -> E::V {
236        self.0.take().unwrap_or_else(|| unreachable!()).remove()
237    }
238}
239
240impl<'a, E> Deref for DefaultOccupiedEntry<'a, E>
241where
242    E: OccupiedEntry<'a, V: Default + Eq>,
243{
244    type Target = E::V;
245
246    #[inline]
247    fn deref(&self) -> &Self::Target {
248        self.0.as_ref().unwrap_or_else(|| unreachable!()).get()
249    }
250}
251
252impl<'a, E> DerefMut for DefaultOccupiedEntry<'a, E>
253where
254    E: OccupiedEntry<'a, V: Default + Eq>,
255{
256    #[inline]
257    fn deref_mut(&mut self) -> &mut Self::Target {
258        self.0.as_mut().unwrap_or_else(|| unreachable!()).get_mut()
259    }
260}
261
262impl<'a, E> Drop for DefaultOccupiedEntry<'a, E>
263where
264    E: OccupiedEntry<'a, V: Default + Eq>,
265{
266    #[inline]
267    fn drop(&mut self) {
268        let Some(entry) = self.0.take() else {
269            return;
270        };
271        if entry.get() == &Default::default() {
272            entry.remove();
273        }
274    }
275}
defuse_map_utils/cleanup.rs

defuse_map_utils/
cleanup.rs
