rayon_hash/par/
map.rs

1/// Rayon extensions to `HashMap`
2use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelExtend, ParallelIterator};
3use std::hash::{BuildHasher, Hash};
4
5use super::table;
6use crate::HashMap;
7
8pub use self::table::{ParIntoIter, ParIter, ParIterMut};
9pub use self::table::{ParKeys, ParValues, ParValuesMut};
10
11impl<K: Sync, V, S> HashMap<K, V, S> {
12    pub fn par_keys(&self) -> ParKeys<K, V> {
13        self.table.par_keys()
14    }
15}
16
17impl<K, V: Sync, S> HashMap<K, V, S> {
18    pub fn par_values(&self) -> ParValues<K, V> {
19        self.table.par_values()
20    }
21}
22
23impl<K, V: Send, S> HashMap<K, V, S> {
24    pub fn par_values_mut(&mut self) -> ParValuesMut<K, V> {
25        self.table.par_values_mut()
26    }
27}
28
29impl<K, V, S> HashMap<K, V, S>
30where
31    K: Eq + Hash + Sync,
32    V: PartialEq + Sync,
33    S: BuildHasher + Sync,
34{
35    pub fn par_eq(&self, other: &Self) -> bool {
36        self.len() == other.len() && self
37            .into_par_iter()
38            .all(|(key, value)| other.get(key).map_or(false, |v| *value == *v))
39    }
40}
41
42impl<K: Send, V: Send, S> IntoParallelIterator for HashMap<K, V, S> {
43    type Item = (K, V);
44    type Iter = ParIntoIter<K, V>;
45
46    fn into_par_iter(self) -> Self::Iter {
47        self.table.into_par_iter()
48    }
49}
50
51impl<'a, K: Sync, V: Sync, S> IntoParallelIterator for &'a HashMap<K, V, S> {
52    type Item = (&'a K, &'a V);
53    type Iter = ParIter<'a, K, V>;
54
55    fn into_par_iter(self) -> Self::Iter {
56        self.table.into_par_iter()
57    }
58}
59
60impl<'a, K: Sync, V: Send, S> IntoParallelIterator for &'a mut HashMap<K, V, S> {
61    type Item = (&'a K, &'a mut V);
62    type Iter = ParIterMut<'a, K, V>;
63
64    fn into_par_iter(self) -> Self::Iter {
65        self.table.into_par_iter()
66    }
67}
68
69/// Collect (key, value) pairs from a parallel iterator into a
70/// hashmap. If multiple pairs correspond to the same key, then the
71/// ones produced earlier in the parallel iterator will be
72/// overwritten, just as with a sequential iterator.
73impl<K, V, S> FromParallelIterator<(K, V)> for HashMap<K, V, S>
74where
75    K: Eq + Hash + Send,
76    V: Send,
77    S: BuildHasher + Default + Send,
78{
79    fn from_par_iter<P>(par_iter: P) -> Self
80    where
81        P: IntoParallelIterator<Item = (K, V)>,
82    {
83        let mut map = HashMap::default();
84        map.par_extend(par_iter);
85        map
86    }
87}
88
89/// Extend a hash map with items from a parallel iterator.
90impl<K, V, S> ParallelExtend<(K, V)> for HashMap<K, V, S>
91where
92    K: Eq + Hash + Send,
93    V: Send,
94    S: BuildHasher + Send,
95{
96    fn par_extend<I>(&mut self, par_iter: I)
97    where
98        I: IntoParallelIterator<Item = (K, V)>,
99    {
100        extend(self, par_iter);
101    }
102}
103
104/// Extend a hash map with copied items from a parallel iterator.
105impl<'a, K, V, S> ParallelExtend<(&'a K, &'a V)> for HashMap<K, V, S>
106where
107    K: Copy + Eq + Hash + Send + Sync,
108    V: Copy + Send + Sync,
109    S: BuildHasher + Send,
110{
111    fn par_extend<I>(&mut self, par_iter: I)
112    where
113        I: IntoParallelIterator<Item = (&'a K, &'a V)>,
114    {
115        extend(self, par_iter);
116    }
117}
118
119// This is equal to the normal `HashMap` -- no custom advantage.
120fn extend<K, V, S, I>(map: &mut HashMap<K, V, S>, par_iter: I)
121where
122    K: Eq + Hash,
123    S: BuildHasher,
124    I: IntoParallelIterator,
125    HashMap<K, V, S>: Extend<I::Item>,
126{
127    let (list, len) = super::collect(par_iter);
128
129    // Keys may be already present or show multiple times in the iterator.
130    // Reserve the entire length if the map is empty.
131    // Otherwise reserve half the length (rounded up), so the map
132    // will only resize twice in the worst case.
133    let reserve = if map.is_empty() { len } else { (len + 1) / 2 };
134    map.reserve(reserve);
135    for vec in list {
136        map.extend(vec);
137    }
138}
139
140#[cfg(test)]
141mod test_par_map {
142    use super::HashMap;
143    use rayon::prelude::*;
144    use std::hash::{Hash, Hasher};
145    use std::sync::atomic::{AtomicUsize, Ordering};
146
147    struct Dropable<'a> {
148        k: usize,
149        counter: &'a AtomicUsize,
150    }
151
152    impl<'a> Dropable<'a> {
153        fn new(k: usize, counter: &AtomicUsize) -> Dropable {
154            counter.fetch_add(1, Ordering::Relaxed);
155
156            Dropable {
157                k: k,
158                counter: counter,
159            }
160        }
161    }
162
163    impl<'a> Drop for Dropable<'a> {
164        fn drop(&mut self) {
165            self.counter.fetch_sub(1, Ordering::Relaxed);
166        }
167    }
168
169    impl<'a> Clone for Dropable<'a> {
170        fn clone(&self) -> Dropable<'a> {
171            Dropable::new(self.k, self.counter)
172        }
173    }
174
175    impl<'a> Hash for Dropable<'a> {
176        fn hash<H>(&self, state: &mut H)
177        where
178            H: Hasher,
179        {
180            self.k.hash(state)
181        }
182    }
183
184    impl<'a> PartialEq for Dropable<'a> {
185        fn eq(&self, other: &Self) -> bool {
186            self.k == other.k
187        }
188    }
189
190    impl<'a> Eq for Dropable<'a> {}
191
192    #[test]
193    fn test_into_iter_drops() {
194        let key = AtomicUsize::new(0);
195        let value = AtomicUsize::new(0);
196
197        let hm = {
198            let mut hm = HashMap::new();
199
200            assert_eq!(key.load(Ordering::Relaxed), 0);
201            assert_eq!(value.load(Ordering::Relaxed), 0);
202
203            for i in 0..100 {
204                let d1 = Dropable::new(i, &key);
205                let d2 = Dropable::new(i + 100, &value);
206                hm.insert(d1, d2);
207            }
208
209            assert_eq!(key.load(Ordering::Relaxed), 100);
210            assert_eq!(value.load(Ordering::Relaxed), 100);
211
212            hm
213        };
214
215        // By the way, ensure that cloning doesn't screw up the dropping.
216        drop(hm.clone());
217
218        {
219            assert_eq!(key.load(Ordering::Relaxed), 100);
220            assert_eq!(value.load(Ordering::Relaxed), 100);
221
222            // retain only half
223            let _v: Vec<_> = hm
224                .into_par_iter()
225                .filter(|&(ref key, _)| key.k < 50)
226                .collect();
227
228            assert_eq!(key.load(Ordering::Relaxed), 50);
229            assert_eq!(value.load(Ordering::Relaxed), 50);
230        };
231
232        assert_eq!(key.load(Ordering::Relaxed), 0);
233        assert_eq!(value.load(Ordering::Relaxed), 0);
234    }
235
236    #[test]
237    fn test_empty_iter() {
238        let mut m: HashMap<isize, bool> = HashMap::new();
239        //assert_eq!(m.par_drain().count(), 0);
240        assert_eq!(m.par_keys().count(), 0);
241        assert_eq!(m.par_values().count(), 0);
242        assert_eq!(m.par_values_mut().count(), 0);
243        assert_eq!(m.par_iter().count(), 0);
244        assert_eq!(m.par_iter_mut().count(), 0);
245        assert_eq!(m.len(), 0);
246        assert!(m.is_empty());
247        assert_eq!(m.into_par_iter().count(), 0);
248    }
249
250    #[test]
251    fn test_iterate() {
252        let mut m = HashMap::with_capacity(4);
253        for i in 0..32 {
254            assert!(m.insert(i, i * 2).is_none());
255        }
256        assert_eq!(m.len(), 32);
257
258        let observed = AtomicUsize::new(0);
259
260        m.par_iter().for_each(|(k, v)| {
261            assert_eq!(*v, *k * 2);
262            observed.fetch_or(1 << *k, Ordering::Relaxed);
263        });
264        assert_eq!(observed.into_inner(), 0xFFFF_FFFF);
265    }
266
267    #[test]
268    fn test_keys() {
269        let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];
270        let map: HashMap<_, _> = vec.into_par_iter().collect();
271        let keys: Vec<_> = map.par_keys().cloned().collect();
272        assert_eq!(keys.len(), 3);
273        assert!(keys.contains(&1));
274        assert!(keys.contains(&2));
275        assert!(keys.contains(&3));
276    }
277
278    #[test]
279    fn test_values() {
280        let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];
281        let map: HashMap<_, _> = vec.into_par_iter().collect();
282        let values: Vec<_> = map.par_values().cloned().collect();
283        assert_eq!(values.len(), 3);
284        assert!(values.contains(&'a'));
285        assert!(values.contains(&'b'));
286        assert!(values.contains(&'c'));
287    }
288
289    #[test]
290    fn test_values_mut() {
291        let vec = vec![(1, 1), (2, 2), (3, 3)];
292        let mut map: HashMap<_, _> = vec.into_par_iter().collect();
293        map.par_values_mut().for_each(|value| *value = (*value) * 2);
294        let values: Vec<_> = map.par_values().cloned().collect();
295        assert_eq!(values.len(), 3);
296        assert!(values.contains(&2));
297        assert!(values.contains(&4));
298        assert!(values.contains(&6));
299    }
300
301    #[test]
302    fn test_eq() {
303        let mut m1 = HashMap::new();
304        m1.insert(1, 2);
305        m1.insert(2, 3);
306        m1.insert(3, 4);
307
308        let mut m2 = HashMap::new();
309        m2.insert(1, 2);
310        m2.insert(2, 3);
311
312        assert!(!m1.par_eq(&m2));
313
314        m2.insert(3, 4);
315
316        assert!(m1.par_eq(&m2));
317    }
318
319    #[test]
320    fn test_from_iter() {
321        let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
322
323        let map: HashMap<_, _> = xs.par_iter().cloned().collect();
324
325        for &(k, v) in &xs {
326            assert_eq!(map.get(&k), Some(&v));
327        }
328    }
329
330    #[test]
331    fn test_extend_ref() {
332        let mut a = HashMap::new();
333        a.insert(1, "one");
334        let mut b = HashMap::new();
335        b.insert(2, "two");
336        b.insert(3, "three");
337
338        a.par_extend(&b);
339
340        assert_eq!(a.len(), 3);
341        assert_eq!(a[&1], "one");
342        assert_eq!(a[&2], "two");
343        assert_eq!(a[&3], "three");
344    }
345}