Skip to main content

cairo_lang_utils/
unordered_hash_map.rs

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