1 use super::mapref::multiple::{RefMulti, RefMutMulti};
2 use super::util;
3 use crate::lock::{RwLockReadGuard, RwLockWriteGuard};
4 use crate::t::Map;
5 use crate::util::SharedValue;
6 use crate::{DashMap, HashMap};
7 use core::hash::{BuildHasher, Hash};
8 use core::mem;
9 use hashbrown::hash_map;
10 use std::collections::hash_map::RandomState;
11 use std::sync::Arc;
12 
13 /// Iterator over a DashMap yielding key value pairs.
14 ///
15 /// # Examples
16 ///
17 /// ```
18 /// use dashmap::DashMap;
19 ///
20 /// let map = DashMap::new();
21 /// map.insert("hello", "world");
22 /// map.insert("alex", "steve");
23 /// let pairs: Vec<(&'static str, &'static str)> = map.into_iter().collect();
24 /// assert_eq!(pairs.len(), 2);
25 /// ```
26 pub struct OwningIter<K, V, S = RandomState> {
27     map: DashMap<K, V, S>,
28     shard_i: usize,
29     current: Option<GuardOwningIter<K, V>>,
30 }
31 
32 impl<K: Eq + Hash, V, S: BuildHasher + Clone> OwningIter<K, V, S> {
new(map: DashMap<K, V, S>) -> Self33     pub(crate) fn new(map: DashMap<K, V, S>) -> Self {
34         Self {
35             map,
36             shard_i: 0,
37             current: None,
38         }
39     }
40 }
41 
42 type GuardOwningIter<K, V> = hash_map::IntoIter<K, SharedValue<V>>;
43 
44 impl<K: Eq + Hash, V, S: BuildHasher + Clone> Iterator for OwningIter<K, V, S> {
45     type Item = (K, V);
46 
next(&mut self) -> Option<Self::Item>47     fn next(&mut self) -> Option<Self::Item> {
48         loop {
49             if let Some(current) = self.current.as_mut() {
50                 if let Some((k, v)) = current.next() {
51                     return Some((k, v.into_inner()));
52                 }
53             }
54 
55             if self.shard_i == self.map._shard_count() {
56                 return None;
57             }
58 
59             //let guard = unsafe { self.map._yield_read_shard(self.shard_i) };
60             let mut shard_wl = unsafe { self.map._yield_write_shard(self.shard_i) };
61 
62             let hasher = self.map._hasher();
63 
64             let map = mem::replace(&mut *shard_wl, HashMap::with_hasher(hasher));
65 
66             drop(shard_wl);
67 
68             let iter = map.into_iter();
69 
70             //unsafe { ptr::write(&mut self.current, Some((arcee, iter))); }
71             self.current = Some(iter);
72 
73             self.shard_i += 1;
74         }
75     }
76 }
77 
78 unsafe impl<K, V, S> Send for OwningIter<K, V, S>
79 where
80     K: Eq + Hash + Send,
81     V: Send,
82     S: BuildHasher + Clone + Send,
83 {
84 }
85 
86 unsafe impl<K, V, S> Sync for OwningIter<K, V, S>
87 where
88     K: Eq + Hash + Sync,
89     V: Sync,
90     S: BuildHasher + Clone + Sync,
91 {
92 }
93 
94 type GuardIter<'a, K, V, S> = (
95     Arc<RwLockReadGuard<'a, HashMap<K, V, S>>>,
96     hash_map::Iter<'a, K, SharedValue<V>>,
97 );
98 
99 type GuardIterMut<'a, K, V, S> = (
100     Arc<RwLockWriteGuard<'a, HashMap<K, V, S>>>,
101     hash_map::IterMut<'a, K, SharedValue<V>>,
102 );
103 
104 /// Iterator over a DashMap yielding immutable references.
105 ///
106 /// # Examples
107 ///
108 /// ```
109 /// use dashmap::DashMap;
110 ///
111 /// let map = DashMap::new();
112 /// map.insert("hello", "world");
113 /// assert_eq!(map.iter().count(), 1);
114 /// ```
115 pub struct Iter<'a, K, V, S = RandomState, M = DashMap<K, V, S>> {
116     map: &'a M,
117     shard_i: usize,
118     current: Option<GuardIter<'a, K, V, S>>,
119 }
120 
121 impl<'i, K: Clone + Hash + Eq, V: Clone, S: Clone + BuildHasher> Clone for Iter<'i, K, V, S> {
clone(&self) -> Self122     fn clone(&self) -> Self {
123         Iter::new(self.map)
124     }
125 }
126 
127 unsafe impl<'a, 'i, K, V, S, M> Send for Iter<'i, K, V, S, M>
128 where
129     K: 'a + Eq + Hash + Send,
130     V: 'a + Send,
131     S: 'a + BuildHasher + Clone,
132     M: Map<'a, K, V, S>,
133 {
134 }
135 
136 unsafe impl<'a, 'i, K, V, S, M> Sync for Iter<'i, K, V, S, M>
137 where
138     K: 'a + Eq + Hash + Sync,
139     V: 'a + Sync,
140     S: 'a + BuildHasher + Clone,
141     M: Map<'a, K, V, S>,
142 {
143 }
144 
145 impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>> Iter<'a, K, V, S, M> {
new(map: &'a M) -> Self146     pub(crate) fn new(map: &'a M) -> Self {
147         Self {
148             map,
149             shard_i: 0,
150             current: None,
151         }
152     }
153 }
154 
155 impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>> Iterator
156     for Iter<'a, K, V, S, M>
157 {
158     type Item = RefMulti<'a, K, V, S>;
159 
next(&mut self) -> Option<Self::Item>160     fn next(&mut self) -> Option<Self::Item> {
161         loop {
162             if let Some(current) = self.current.as_mut() {
163                 if let Some((k, v)) = current.1.next() {
164                     let guard = current.0.clone();
165 
166                     return unsafe { Some(RefMulti::new(guard, k, v.get())) };
167                 }
168             }
169 
170             if self.shard_i == self.map._shard_count() {
171                 return None;
172             }
173 
174             let guard = unsafe { self.map._yield_read_shard(self.shard_i) };
175 
176             let sref: &HashMap<K, V, S> = unsafe { util::change_lifetime_const(&*guard) };
177 
178             let iter = sref.iter();
179 
180             self.current = Some((Arc::new(guard), iter));
181 
182             self.shard_i += 1;
183         }
184     }
185 }
186 
187 /// Iterator over a DashMap yielding mutable references.
188 ///
189 /// # Examples
190 ///
191 /// ```
192 /// use dashmap::DashMap;
193 ///
194 /// let map = DashMap::new();
195 /// map.insert("Johnny", 21);
196 /// map.iter_mut().for_each(|mut r| *r += 1);
197 /// assert_eq!(*map.get("Johnny").unwrap(), 22);
198 /// ```
199 pub struct IterMut<'a, K, V, S = RandomState, M = DashMap<K, V, S>> {
200     map: &'a M,
201     shard_i: usize,
202     current: Option<GuardIterMut<'a, K, V, S>>,
203 }
204 
205 unsafe impl<'a, 'i, K, V, S, M> Send for IterMut<'i, K, V, S, M>
206 where
207     K: 'a + Eq + Hash + Send,
208     V: 'a + Send,
209     S: 'a + BuildHasher + Clone,
210     M: Map<'a, K, V, S>,
211 {
212 }
213 
214 unsafe impl<'a, 'i, K, V, S, M> Sync for IterMut<'i, K, V, S, M>
215 where
216     K: 'a + Eq + Hash + Sync,
217     V: 'a + Sync,
218     S: 'a + BuildHasher + Clone,
219     M: Map<'a, K, V, S>,
220 {
221 }
222 
223 impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>>
224     IterMut<'a, K, V, S, M>
225 {
new(map: &'a M) -> Self226     pub(crate) fn new(map: &'a M) -> Self {
227         Self {
228             map,
229             shard_i: 0,
230             current: None,
231         }
232     }
233 }
234 
235 impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>> Iterator
236     for IterMut<'a, K, V, S, M>
237 {
238     type Item = RefMutMulti<'a, K, V, S>;
239 
next(&mut self) -> Option<Self::Item>240     fn next(&mut self) -> Option<Self::Item> {
241         loop {
242             if let Some(current) = self.current.as_mut() {
243                 if let Some((k, v)) = current.1.next() {
244                     let guard = current.0.clone();
245 
246                     unsafe {
247                         let k = util::change_lifetime_const(k);
248 
249                         let v = &mut *v.as_ptr();
250 
251                         return Some(RefMutMulti::new(guard, k, v));
252                     }
253                 }
254             }
255 
256             if self.shard_i == self.map._shard_count() {
257                 return None;
258             }
259 
260             let mut guard = unsafe { self.map._yield_write_shard(self.shard_i) };
261 
262             let sref: &mut HashMap<K, V, S> = unsafe { util::change_lifetime_mut(&mut *guard) };
263 
264             let iter = sref.iter_mut();
265 
266             self.current = Some((Arc::new(guard), iter));
267 
268             self.shard_i += 1;
269         }
270     }
271 }
272 
273 #[cfg(test)]
274 mod tests {
275     use crate::DashMap;
276 
277     #[test]
iter_mut_manual_count()278     fn iter_mut_manual_count() {
279         let map = DashMap::new();
280 
281         map.insert("Johnny", 21);
282 
283         assert_eq!(map.len(), 1);
284 
285         let mut c = 0;
286 
287         for shard in map.shards() {
288             c += shard.write().iter_mut().count();
289         }
290 
291         assert_eq!(c, 1);
292     }
293 
294     #[test]
iter_mut_count()295     fn iter_mut_count() {
296         let map = DashMap::new();
297 
298         map.insert("Johnny", 21);
299 
300         assert_eq!(map.len(), 1);
301 
302         assert_eq!(map.iter_mut().count(), 1);
303     }
304 
305     #[test]
iter_count()306     fn iter_count() {
307         let map = DashMap::new();
308 
309         map.insert("Johnny", 21);
310 
311         assert_eq!(map.len(), 1);
312 
313         assert_eq!(map.iter().count(), 1);
314     }
315 }
316