cairo_lang_utils/
unordered_hash_map.rs

1#[cfg(test)]
2#[path = "unordered_hash_map_test.rs"]
3mod test;
4
5#[cfg(not(feature = "std"))]
6use alloc::vec;
7use core::borrow::Borrow;
8use core::hash::{BuildHasher, Hash};
9use core::ops::Index;
10#[cfg(feature = "std")]
11use std::collections::HashMap;
12#[cfg(feature = "std")]
13pub use std::collections::hash_map::Entry;
14#[cfg(feature = "std")]
15use std::collections::hash_map::OccupiedEntry;
16#[cfg(feature = "std")]
17use std::collections::hash_map::RandomState;
18#[cfg(feature = "std")]
19use std::vec;
20
21#[cfg(not(feature = "std"))]
22use hashbrown::HashMap;
23#[cfg(not(feature = "std"))]
24pub use hashbrown::hash_map::Entry;
25use itertools::Itertools;
26
27#[cfg(feature = "std")]
28type BHImpl = RandomState;
29#[cfg(not(feature = "std"))]
30type BHImpl = hashbrown::DefaultHashBuilder;
31
32/// A hash map that does not care about the order of insertion.
33///
34/// In particular, it does not support iterating, in order to guarantee deterministic compilation.
35/// It does support aggregation which can be used in intermediate computations (see `aggregate_by`).
36/// For an iterable version see [OrderedHashMap](crate::ordered_hash_map::OrderedHashMap).
37#[derive(Clone, Debug)]
38pub struct UnorderedHashMap<Key, Value, BH = BHImpl>(HashMap<Key, Value, BH>);
39
40#[cfg(feature = "salsa")]
41unsafe impl<Key: salsa::Update + Eq + Hash, Value: salsa::Update> salsa::Update
42    for UnorderedHashMap<Key, Value, BHImpl>
43{
44    // This code was taken from the salsa::Update trait implementation for IndexMap.
45    // It is defined privately in macro_rules! maybe_update_map in the db-ext-macro repo.
46    unsafe fn maybe_update(old_pointer: *mut Self, new_map: Self) -> bool {
47        let old_map: &mut Self = unsafe { &mut *old_pointer };
48
49        // To be considered "equal", the set of keys
50        // must be the same between the two maps.
51        let same_keys =
52            old_map.len() == new_map.len() && old_map.0.keys().all(|k| new_map.0.contains_key(k));
53
54        // If the set of keys has changed, then just pull in the new values
55        // from new_map and discard the old ones.
56        if !same_keys {
57            old_map.0.clear();
58            old_map.0.extend(new_map.0);
59            return true;
60        }
61
62        // Otherwise, recursively descend to the values.
63        // We do not invoke `K::update` because we assume
64        // that if the values are `Eq` they must not need
65        // updating (see the trait criteria).
66        let mut changed = false;
67        for (key, new_value) in new_map.0.into_iter() {
68            let old_value = old_map.0.get_mut(&key).unwrap();
69            changed |= unsafe { Value::maybe_update(old_value, new_value) };
70        }
71        changed
72    }
73}
74
75impl<Key, Value, BH> UnorderedHashMap<Key, Value, BH> {
76    fn with_hasher(hash_builder: BH) -> Self {
77        Self(HashMap::<Key, Value, BH>::with_hasher(hash_builder))
78    }
79}
80
81impl<Key, Value, BH> PartialEq for UnorderedHashMap<Key, Value, BH>
82where
83    Key: Eq + Hash,
84    Value: PartialEq,
85    BH: BuildHasher,
86{
87    fn eq(&self, other: &Self) -> bool {
88        self.0 == other.0
89    }
90}
91
92impl<Key, Value, BH> Eq for UnorderedHashMap<Key, Value, BH>
93where
94    Key: Eq + Hash,
95    Value: Eq,
96    BH: BuildHasher,
97{
98}
99
100impl<Key, Value, BH> UnorderedHashMap<Key, Value, BH> {
101    /// Returns the number of elements in the map.
102    pub fn len(&self) -> usize {
103        self.0.len()
104    }
105
106    /// Returns true if the map contains no elements.
107    pub fn is_empty(&self) -> bool {
108        self.0.is_empty()
109    }
110}
111
112impl<Key: Eq + Hash, Value, BH: BuildHasher> UnorderedHashMap<Key, Value, BH> {
113    /// Returns a reference to the value corresponding to the key.
114    ///
115    /// The key may be any borrowed form of the map's key type, but [`Hash`] and [`Eq`] on the
116    /// borrowed form *must* match those for the key type.
117    pub fn get<Q>(&self, key: &Q) -> Option<&Value>
118    where
119        Key: Borrow<Q>,
120        Q: Hash + Eq + ?Sized,
121    {
122        self.0.get(key)
123    }
124
125    /// Returns a mutable reference to the value corresponding to the key.
126    ///
127    /// The key may be any borrowed form of the map's key type, but [`Hash`] and [`Eq`] on the
128    /// borrowed form *must* match those for the key type.
129    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut Value>
130    where
131        Key: Borrow<Q>,
132        Q: Hash + Eq + ?Sized,
133    {
134        self.0.get_mut(key)
135    }
136
137    /// Inserts a key-value pair into the map.
138    ///
139    /// If the map did not have this key present, None is returned.
140    ///
141    /// If the map did have this key present, the value is updated, and the old value is returned.
142    /// The key is not updated, though; this matters for types that can be == without being
143    /// identical.
144    pub fn insert(&mut self, key: Key, value: Value) -> Option<Value> {
145        self.0.insert(key, value)
146    }
147
148    /// Removes a key from the map, returning the value at the key if the key was previously in the
149    /// map.
150    ///
151    /// The key may be any borrowed form of the map's key type, but Hash and Eq on the borrowed form
152    /// must match those for the key type.
153    pub fn remove<Q>(&mut self, key: &Q) -> Option<Value>
154    where
155        Key: Borrow<Q>,
156        Q: Hash + Eq + ?Sized,
157    {
158        self.0.remove(key)
159    }
160
161    #[cfg(feature = "std")]
162    /// Gets the given key's corresponding entry in the map for in-place manipulation.
163    pub fn entry(&mut self, key: Key) -> Entry<'_, Key, Value> {
164        self.0.entry(key)
165    }
166
167    #[cfg(not(feature = "std"))]
168    /// Gets the given key's corresponding entry in the map for in-place manipulation.
169    pub fn entry(&mut self, key: Key) -> Entry<'_, Key, Value, BH> {
170        self.0.entry(key)
171    }
172
173    /// Returns true if the map contains a value for the specified key.
174    pub fn contains_key<Q>(&self, key: &Q) -> bool
175    where
176        Q: ?Sized,
177        Key: Borrow<Q>,
178        Q: Hash + Eq,
179    {
180        self.0.contains_key(key)
181    }
182
183    /// Maps the values of the map to new values using the given function.
184    pub fn map<TargetValue>(
185        &self,
186        mapper: impl Fn(&Value) -> TargetValue,
187    ) -> UnorderedHashMap<Key, TargetValue, BH>
188    where
189        Key: Clone,
190        BH: Clone,
191    {
192        self.0.iter().fold(
193            UnorderedHashMap::<_, _, _>::with_hasher(self.0.hasher().clone()),
194            |mut acc, (key, value)| {
195                match acc.entry(key.clone()) {
196                    Entry::Occupied(_) => {
197                        unreachable!("The original map should not contain duplicate keys.");
198                    }
199                    Entry::Vacant(vacant) => {
200                        vacant.insert(mapper(value));
201                    }
202                };
203                acc
204            },
205        )
206    }
207
208    /// Aggregates values of the map using the given functions.
209    /// `mapping_function` maps each key to a new key, possibly mapping multiple original keys to
210    /// the same target key.
211    /// `reduce_function` dictates how to aggregate any two values of the same target key.
212    /// `default_value` is the initial value for each target key.
213    /// Note! as the map is unordered, `reduce_function` should be commutative. Otherwise, the
214    /// result is undefined (nondeterministic).
215    pub fn aggregate_by<TargetKey: Eq + Hash, TargetValue>(
216        &self,
217        mapping_function: impl Fn(&Key) -> TargetKey,
218        reduce_function: impl Fn(&TargetValue, &Value) -> TargetValue,
219        default_value: &TargetValue,
220    ) -> UnorderedHashMap<TargetKey, TargetValue, BH>
221    where
222        BH: Clone,
223    {
224        self.0.iter().fold(
225            UnorderedHashMap::<_, _, _>::with_hasher(self.0.hasher().clone()),
226            |mut acc, (key, value)| {
227                let target_key = mapping_function(key);
228                match acc.entry(target_key) {
229                    Entry::Occupied(occupied) => {
230                        let old_target_value = occupied.into_mut();
231                        let new_target_value = reduce_function(old_target_value, value);
232                        *old_target_value = new_target_value;
233                    }
234                    Entry::Vacant(vacant) => {
235                        let new_value = reduce_function(default_value, value);
236                        vacant.insert(new_value);
237                    }
238                };
239                acc
240            },
241        )
242    }
243
244    /// Iterates the map in an ascending order applied by the `Ord` implementation of `Key`.
245    /// NOTE! To guarantee a deterministic output, the `Ord` implementation must apply a strict
246    /// ordering. That is, `a <= b` and `b <= a`, then `a == b`. If `Ord` is derived (in all
247    /// hierarchy levels), this is probably the case. If the ordering is not strict, the order of
248    /// the output OrderedHashMap is undefined (nondeterministic).
249    /// This can be used to convert an unordered map to an ordered map (mostly when the unordered
250    /// map was used for intermediate processing).
251    pub fn iter_sorted(&self) -> impl Iterator<Item = (&Key, &Value)>
252    where
253        Key: Ord,
254    {
255        self.0.iter().sorted_by_key(|(key, _)| *key)
256    }
257
258    /// A consuming version of `iter_sorted`.
259    pub fn into_iter_sorted(self) -> impl Iterator<Item = (Key, Value)>
260    where
261        Key: Ord,
262    {
263        self.0.into_iter().sorted_by(|(lhs, _), (rhs, _)| lhs.cmp(rhs))
264    }
265
266    /// Iterates the map in an ascending order of the keys produced by the given function `f`.
267    /// NOTE! To guarantee a deterministic output, `f`'s implementation must apply a strict
268    /// ordering of the (Key, Value) pairs. That is, for any given pair of entries `a=(k_a, v_a)`
269    /// and `b=(k_b, v_b)`, if `a <= b` and `b <= a`, then `a == b`. If the ordering is not strict,
270    /// the order of the output OrderedHashMap is undefined (nondeterministic).
271    /// This can be used to convert an unordered map to an ordered map (mostly when the unordered
272    /// map was used for intermediate processing).
273    pub fn iter_sorted_by_key<TargetKey, F>(&self, f: F) -> vec::IntoIter<(&Key, &Value)>
274    where
275        TargetKey: Ord,
276        F: FnMut(&(&Key, &Value)) -> TargetKey,
277    {
278        self.0.iter().sorted_by_key(f)
279    }
280
281    /// A consuming version of `iter_sorted_by_key`.
282    pub fn into_iter_sorted_by_key<TargetKey, F>(self, f: F) -> vec::IntoIter<(Key, Value)>
283    where
284        TargetKey: Ord,
285        F: FnMut(&(Key, Value)) -> TargetKey,
286    {
287        self.0.into_iter().sorted_by_key(f)
288    }
289
290    /// Creates a new map with only the elements from the original map for which the given predicate
291    /// returns `true`. Consuming.
292    pub fn filter<P>(self, mut p: P) -> Self
293    where
294        BH: Default,
295        P: FnMut(&Key, &Value) -> bool,
296    {
297        Self(self.0.into_iter().filter(|(key, value)| p(key, value)).collect())
298    }
299
300    /// Non consuming version of `filter`. Only clones the filtered entries. Requires `Key` and
301    /// `Value` to implement `Clone`.
302    pub fn filter_cloned<P>(&self, mut p: P) -> Self
303    where
304        BH: Default,
305        P: FnMut(&Key, &Value) -> bool,
306        Key: Clone,
307        Value: Clone,
308    {
309        Self(
310            self.0
311                .iter()
312                .filter_map(
313                    |(key, value)| {
314                        if p(key, value) { Some((key.clone(), value.clone())) } else { None }
315                    },
316                )
317                .collect(),
318        )
319    }
320
321    #[cfg(feature = "std")]
322    /// Merges the map with another map. If a key is present in both maps, the given handler
323    /// function is used to combine the values.
324    pub fn merge<HandleDuplicate>(&mut self, other: &Self, handle_duplicate: HandleDuplicate)
325    where
326        BH: Clone,
327        HandleDuplicate: Fn(OccupiedEntry<'_, Key, Value>, &Value),
328        Key: Clone,
329        Value: Clone,
330    {
331        for (key, value) in &other.0 {
332            match self.0.entry(key.clone()) {
333                Entry::Occupied(e) => {
334                    handle_duplicate(e, value);
335                }
336                Entry::Vacant(e) => {
337                    e.insert(value.clone());
338                }
339            }
340        }
341    }
342
343    /// Clears the map, removing all key-value pairs. Keeps the allocated memory for reuse.
344    pub fn clear(&mut self) {
345        self.0.clear();
346    }
347}
348
349impl<Key, Q: ?Sized, Value, BH: BuildHasher> Index<&Q> for UnorderedHashMap<Key, Value, BH>
350where
351    Key: Eq + Hash + Borrow<Q>,
352    Q: Eq + Hash,
353{
354    type Output = Value;
355
356    fn index(&self, key: &Q) -> &Self::Output {
357        self.0.index(key)
358    }
359}
360
361impl<Key, Value, BH: Default> Default for UnorderedHashMap<Key, Value, BH> {
362    #[cfg(feature = "std")]
363    fn default() -> Self {
364        Self(Default::default())
365    }
366    #[cfg(not(feature = "std"))]
367    fn default() -> Self {
368        Self(HashMap::with_hasher(Default::default()))
369    }
370}
371
372impl<Key: Hash + Eq, Value, BH: BuildHasher + Default> FromIterator<(Key, Value)>
373    for UnorderedHashMap<Key, Value, BH>
374{
375    fn from_iter<T: IntoIterator<Item = (Key, Value)>>(iter: T) -> Self {
376        Self(iter.into_iter().collect())
377    }
378}
379
380impl<Key: Hash + Eq, Value, const N: usize, BH: BuildHasher + Default> From<[(Key, Value); N]>
381    for UnorderedHashMap<Key, Value, BH>
382{
383    fn from(items: [(Key, Value); N]) -> Self {
384        Self(HashMap::from_iter(items))
385    }
386}
387
388impl<Key: Hash + Eq, Value, BH: BuildHasher> Extend<(Key, Value)>
389    for UnorderedHashMap<Key, Value, BH>
390{
391    fn extend<T: IntoIterator<Item = (Key, Value)>>(&mut self, iter: T) {
392        self.0.extend(iter)
393    }
394}