erg_common/
dict.rs

1use std::borrow::Borrow;
2use std::collections::hash_map::{
3    Entry, IntoKeys, IntoValues, Iter, IterMut, Keys, Values, ValuesMut,
4};
5use std::fmt::{self, Write};
6use std::hash::{Hash, Hasher};
7use std::iter::FromIterator;
8use std::ops::{Index, IndexMut};
9
10use crate::fxhash::FxHashMap;
11use crate::get_hash;
12use crate::traits::Immutable;
13
14#[macro_export]
15macro_rules! dict {
16    () => { $crate::dict::Dict::new() };
17    ($($k: expr => $v: expr),+ $(,)?) => {{
18        let mut dict = $crate::dict::Dict::new();
19        $(dict.insert($k, $v);)+
20        dict
21    }};
22}
23
24#[derive(Debug, Clone)]
25pub struct Dict<K, V> {
26    dict: FxHashMap<K, V>,
27}
28
29impl<K: Hash + Eq + Immutable, V: Hash + Eq> PartialEq for Dict<K, V> {
30    fn eq(&self, other: &Dict<K, V>) -> bool {
31        self.dict == other.dict
32    }
33}
34
35impl<K: Hash + Eq + Immutable, V: Hash + Eq> Eq for Dict<K, V> {}
36
37impl<K: Hash, V: Hash> Hash for Dict<K, V> {
38    fn hash<H: Hasher>(&self, state: &mut H) {
39        let len = self.len();
40        len.hash(state);
41        if len <= 1 {
42            for (key, val) in self.iter() {
43                key.hash(state);
44                val.hash(state);
45            }
46        } else {
47            let mut v = self
48                .iter()
49                .map(|(key, val)| (get_hash(key), val))
50                .collect::<Vec<_>>();
51            v.sort_unstable_by_key(|(h, _)| *h);
52            for (h, val) in v.iter() {
53                state.write_usize(*h);
54                val.hash(state);
55            }
56        }
57    }
58}
59
60impl<K: fmt::Display, V: fmt::Display> fmt::Display for Dict<K, V> {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        let mut s = "".to_string();
63        for (k, v) in self.dict.iter() {
64            write!(s, "{k}: {v}, ")?;
65        }
66        s.pop();
67        s.pop();
68        write!(f, "{{{s}}}")
69    }
70}
71
72impl<K: Hash + Eq, V> FromIterator<(K, V)> for Dict<K, V> {
73    #[inline]
74    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Dict<K, V> {
75        let mut dict = Dict::new();
76        dict.extend(iter);
77        dict
78    }
79}
80
81impl<K: Hash + Eq, V> Extend<(K, V)> for Dict<K, V> {
82    #[inline]
83    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
84        self.guaranteed_extend(iter);
85    }
86}
87
88impl<K: Hash + Eq, V> From<Vec<(K, V)>> for Dict<K, V> {
89    #[inline]
90    fn from(v: Vec<(K, V)>) -> Dict<K, V> {
91        v.into_iter().collect()
92    }
93}
94
95impl<K: Hash + Eq + Immutable, V, Q: ?Sized> Index<&Q> for Dict<K, V>
96where
97    K: Borrow<Q>,
98    Q: Hash + Eq,
99{
100    type Output = V;
101    #[inline]
102    fn index(&self, index: &Q) -> &V {
103        self.dict.get(index).unwrap()
104    }
105}
106
107impl<K: Hash + Eq + Immutable, V, Q: ?Sized> IndexMut<&Q> for Dict<K, V>
108where
109    K: Borrow<Q>,
110    Q: Hash + Eq,
111{
112    #[inline]
113    fn index_mut(&mut self, index: &Q) -> &mut V {
114        self.dict.get_mut(index).unwrap()
115    }
116}
117
118impl<K, V> Default for Dict<K, V> {
119    fn default() -> Dict<K, V> {
120        Dict::new()
121    }
122}
123
124impl<K: Clone + Hash + Eq, V: Clone> Dict<&K, &V> {
125    pub fn cloned(&self) -> Dict<K, V> {
126        self.dict
127            .iter()
128            .map(|(&k, &v)| (k.clone(), v.clone()))
129            .collect()
130    }
131}
132
133impl<K, V> Dict<K, V> {
134    #[inline]
135    pub fn new() -> Self {
136        Self {
137            dict: FxHashMap::default(),
138        }
139    }
140
141    /// ```
142    /// # use erg_common::dict;
143    /// # use erg_common::dict::Dict;
144    /// let mut dict = Dict::with_capacity(3);
145    /// assert_eq!(dict.capacity(), 3);
146    /// dict.insert("a", 1);
147    /// assert_eq!(dict.capacity(), 3);
148    /// dict.insert("b", 2);
149    /// dict.insert("c", 3);
150    /// dict.insert("d", 4);
151    /// assert_ne!(dict.capacity(), 3);
152    /// ```
153    pub fn with_capacity(capacity: usize) -> Self {
154        Self {
155            dict: FxHashMap::with_capacity_and_hasher(capacity, Default::default()),
156        }
157    }
158
159    #[inline]
160    pub fn len(&self) -> usize {
161        self.dict.len()
162    }
163
164    #[inline]
165    pub fn is_empty(&self) -> bool {
166        self.dict.is_empty()
167    }
168
169    #[inline]
170    pub fn capacity(&self) -> usize {
171        self.dict.capacity()
172    }
173
174    #[inline]
175    pub fn keys(&self) -> Keys<'_, K, V> {
176        self.dict.keys()
177    }
178
179    #[inline]
180    pub fn values(&self) -> Values<'_, K, V> {
181        self.dict.values()
182    }
183
184    #[inline]
185    pub fn values_mut(&mut self) -> ValuesMut<'_, K, V> {
186        self.dict.values_mut()
187    }
188
189    #[inline]
190    pub fn into_values(self) -> IntoValues<K, V> {
191        self.dict.into_values()
192    }
193
194    #[inline]
195    pub fn into_keys(self) -> IntoKeys<K, V> {
196        self.dict.into_keys()
197    }
198
199    #[inline]
200    pub fn iter(&self) -> Iter<'_, K, V> {
201        self.dict.iter()
202    }
203
204    #[inline]
205    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
206        self.dict.iter_mut()
207    }
208
209    pub fn clear(&mut self) {
210        self.dict.clear();
211    }
212
213    /// remove all elements for which the predicate returns false
214    pub fn retain<F>(&mut self, f: F)
215    where
216        F: FnMut(&K, &mut V) -> bool,
217    {
218        self.dict.retain(f);
219    }
220
221    pub fn retained(mut self, f: impl FnMut(&K, &mut V) -> bool) -> Self {
222        self.retain(f);
223        self
224    }
225
226    pub fn get_by(&self, k: &K, cmp: impl Fn(&K, &K) -> bool) -> Option<&V> {
227        for (k_, v) in self.dict.iter() {
228            if cmp(k, k_) {
229                return Some(v);
230            }
231        }
232        None
233    }
234}
235
236impl<K, V> IntoIterator for Dict<K, V> {
237    type Item = (K, V);
238    type IntoIter = <FxHashMap<K, V> as IntoIterator>::IntoIter;
239    #[inline]
240    fn into_iter(self) -> Self::IntoIter {
241        self.dict.into_iter()
242    }
243}
244
245impl<'a, K, V> IntoIterator for &'a Dict<K, V> {
246    type Item = (&'a K, &'a V);
247    type IntoIter = Iter<'a, K, V>;
248    #[inline]
249    fn into_iter(self) -> Self::IntoIter {
250        self.dict.iter()
251    }
252}
253
254impl<K: Eq, V> Dict<K, V> {
255    /// K: interior-mutable
256    pub fn linear_get<Q>(&self, key: &Q) -> Option<&V>
257    where
258        K: Borrow<Q>,
259        Q: Eq + ?Sized,
260    {
261        self.dict
262            .iter()
263            .find(|(k, _)| (*k).borrow() == key)
264            .map(|(_, v)| v)
265    }
266
267    pub fn linear_get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
268    where
269        K: Borrow<Q>,
270        Q: Eq + ?Sized,
271    {
272        self.dict
273            .iter_mut()
274            .find(|(k, _)| (*k).borrow() == key)
275            .map(|(_, v)| v)
276    }
277}
278
279impl<K: Eq, V: Eq> Dict<K, V> {
280    /// K: interior-mutable
281    pub fn linear_eq(&self, other: &Self) -> bool {
282        if self.len() != other.len() {
283            return false;
284        }
285        for (k, v) in self.iter() {
286            if other.linear_get(k) != Some(v) {
287                return false;
288            }
289        }
290        true
291    }
292}
293
294impl<K: Hash + Eq + Immutable, V> Dict<K, V> {
295    #[inline]
296    pub fn get<Q>(&self, k: &Q) -> Option<&V>
297    where
298        K: Borrow<Q>,
299        Q: Hash + Eq + ?Sized,
300    {
301        self.dict.get(k)
302    }
303
304    #[inline]
305    pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
306    where
307        K: Borrow<Q>,
308        Q: Hash + Eq + ?Sized,
309    {
310        self.dict.get_mut(k)
311    }
312
313    pub fn get_key_value<Q>(&self, k: &Q) -> Option<(&K, &V)>
314    where
315        K: Borrow<Q>,
316        Q: Hash + Eq + ?Sized,
317    {
318        self.dict.get_key_value(k)
319    }
320
321    #[inline]
322    pub fn contains_key<Q>(&self, k: &Q) -> bool
323    where
324        K: Borrow<Q>,
325        Q: Hash + Eq + ?Sized,
326    {
327        self.dict.contains_key(k)
328    }
329
330    #[inline]
331    pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
332    where
333        K: Borrow<Q>,
334        Q: Hash + Eq + ?Sized,
335    {
336        self.dict.remove(k)
337    }
338
339    pub fn remove_entry<Q>(&mut self, k: &Q) -> Option<(K, V)>
340    where
341        K: Borrow<Q>,
342        Q: Hash + Eq + ?Sized,
343    {
344        self.dict.remove_entry(k)
345    }
346
347    pub fn remove_entries<'q, Q>(&mut self, keys: impl IntoIterator<Item = &'q Q>)
348    where
349        K: Borrow<Q>,
350        Q: Hash + Eq + ?Sized + 'q,
351    {
352        for k in keys {
353            self.remove_entry(k);
354        }
355    }
356}
357
358impl<K: Hash + Eq, V> Dict<K, V> {
359    #[inline]
360    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
361        self.dict.insert(k, v)
362    }
363
364    /// NOTE: This method does not consider pairing with values and keys. That is, a value may be paired with a different key (can be considered equal).
365    /// If you need to consider the pairing of the keys and values, use `guaranteed_extend` instead.
366    #[inline]
367    pub fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
368        self.dict.extend(iter);
369    }
370
371    /// If the key already exists, the value will not be updated.
372    #[inline]
373    pub fn guaranteed_extend<I: IntoIterator<Item = (K, V)>>(&mut self, other: I) {
374        for (k, v) in other {
375            self.dict.entry(k).or_insert(v);
376        }
377    }
378
379    #[inline]
380    pub fn merge(&mut self, other: Self) {
381        self.dict.extend(other.dict);
382    }
383
384    #[inline]
385    pub fn concat(mut self, other: Self) -> Self {
386        self.merge(other);
387        self
388    }
389
390    #[inline]
391    pub fn diff(mut self, other: &Self) -> Self {
392        for k in other.dict.keys() {
393            self.dict.remove(k);
394        }
395        self
396    }
397
398    pub fn entry(&mut self, k: K) -> Entry<'_, K, V> {
399        self.dict.entry(k)
400    }
401}
402
403#[cfg(test)]
404mod tests {
405    use super::*;
406
407    use crate::str::Str;
408
409    #[test]
410    fn test_dict() {
411        let mut dict = Dict::new();
412        dict.insert(Str::from("a"), 1);
413        dict.insert(Str::from("b"), 2);
414        dict.insert(Str::from("c"), 3);
415        assert_eq!(dict.len(), 3);
416        assert_eq!(dict.get(&Str::from("a")), Some(&1));
417        assert_eq!(dict.get(&Str::from("b")), Some(&2));
418        assert_eq!(dict.get(&Str::from("c")), Some(&3));
419        assert_eq!(dict.get(&Str::from("d")), None);
420        assert_eq!(dict.get("a"), Some(&1));
421        assert_eq!(dict.get("b"), Some(&2));
422        assert_eq!(dict.get("c"), Some(&3));
423        assert_eq!(dict.get("d"), None);
424        assert_eq!(dict.remove(&Str::from("a")), Some(1));
425        assert_eq!(dict.remove(&Str::from("a")), None);
426        assert_eq!(dict.len(), 2);
427        assert_eq!(dict.get(&Str::from("a")), None);
428        assert_eq!(dict.get(&Str::from("b")), Some(&2));
429        assert_eq!(dict.get(&Str::from("c")), Some(&3));
430        assert_eq!(dict.get(&Str::from("d")), None);
431        dict.clear();
432        assert_eq!(dict.len(), 0);
433        assert_eq!(dict.get(&Str::from("a")), None);
434        assert_eq!(dict.get(&Str::from("b")), None);
435        assert_eq!(dict.get(&Str::from("c")), None);
436        assert_eq!(dict.get(&Str::from("d")), None);
437    }
438}
erg_common/dict.rs

erg_common/
dict.rs
