1use super::mapref::multiple::{RefMulti, RefMutMulti};
2use crate::lock::{RwLockReadGuard, RwLockWriteGuard};
3use crate::t::Map;
4use crate::util::SharedValue;
5use crate::{DashMap, HashMap};
6use core::hash::{BuildHasher, Hash};
7use core::mem;
8use std::collections::hash_map::RandomState;
9use std::marker::PhantomData;
10use std::sync::Arc;
11
12pub struct OwningIter<K, V, S = RandomState> {
26 map: DashMap<K, V, S>,
27 shard_i: usize,
28 current: Option<GuardOwningIter<K, V>>,
29}
30
31impl<K: Eq + Hash, V, S: BuildHasher + Clone> OwningIter<K, V, S> {
32 pub(crate) fn new(map: DashMap<K, V, S>) -> Self {
33 Self {
34 map,
35 shard_i: 0,
36 current: None,
37 }
38 }
39}
40
41type GuardOwningIter<K, V> = hashbrown::raw::RawIntoIter<(K, SharedValue<V>)>;
42
43impl<K: Eq + Hash, V, S: BuildHasher + Clone> Iterator for OwningIter<K, V, S> {
44 type Item = (K, V);
45
46 fn next(&mut self) -> Option<Self::Item> {
47 loop {
48 if let Some(current) = self.current.as_mut() {
49 if let Some((k, v)) = current.next() {
50 return Some((k, v.into_inner()));
51 }
52 }
53
54 if self.shard_i == self.map._shard_count() {
55 return None;
56 }
57
58 let mut shard_wl = unsafe { self.map._yield_write_shard(self.shard_i) };
60
61 let map = mem::take(&mut *shard_wl);
62
63 drop(shard_wl);
64
65 let iter = map.into_iter();
66
67 self.current = Some(iter);
69
70 self.shard_i += 1;
71 }
72 }
73}
74
75unsafe impl<K, V, S> Send for OwningIter<K, V, S>
76where
77 K: Eq + Hash + Send,
78 V: Send,
79 S: BuildHasher + Clone + Send,
80{
81}
82
83unsafe impl<K, V, S> Sync for OwningIter<K, V, S>
84where
85 K: Eq + Hash + Sync,
86 V: Sync,
87 S: BuildHasher + Clone + Sync,
88{
89}
90
91type GuardIter<'a, K, V> = (
92 Arc<RwLockReadGuard<'a, HashMap<K, V>>>,
93 hashbrown::raw::RawIter<(K, SharedValue<V>)>,
94);
95
96type GuardIterMut<'a, K, V> = (
97 Arc<RwLockWriteGuard<'a, HashMap<K, V>>>,
98 hashbrown::raw::RawIter<(K, SharedValue<V>)>,
99);
100
101pub struct Iter<'a, K, V, S = RandomState, M = DashMap<K, V, S>> {
113 map: &'a M,
114 shard_i: usize,
115 current: Option<GuardIter<'a, K, V>>,
116 marker: PhantomData<S>,
117}
118
119impl<'i, K: Clone + Hash + Eq, V: Clone, S: Clone + BuildHasher> Clone for Iter<'i, K, V, S> {
120 fn clone(&self) -> Self {
121 Iter::new(self.map)
122 }
123}
124
125unsafe impl<'a, 'i, K, V, S, M> Send for Iter<'i, K, V, S, M>
126where
127 K: 'a + Eq + Hash + Send,
128 V: 'a + Send,
129 S: 'a + BuildHasher + Clone,
130 M: Map<'a, K, V, S>,
131{
132}
133
134unsafe impl<'a, 'i, K, V, S, M> Sync for Iter<'i, K, V, S, M>
135where
136 K: 'a + Eq + Hash + Sync,
137 V: 'a + Sync,
138 S: 'a + BuildHasher + Clone,
139 M: Map<'a, K, V, S>,
140{
141}
142
143impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>> Iter<'a, K, V, S, M> {
144 pub(crate) fn new(map: &'a M) -> Self {
145 Self {
146 map,
147 shard_i: 0,
148 current: None,
149 marker: PhantomData,
150 }
151 }
152}
153
154impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>> Iterator
155 for Iter<'a, K, V, S, M>
156{
157 type Item = RefMulti<'a, K, V>;
158
159 fn next(&mut self) -> Option<Self::Item> {
160 loop {
161 if let Some(current) = self.current.as_mut() {
162 if let Some(b) = current.1.next() {
163 return unsafe {
164 let (k, v) = b.as_ref();
165 let guard = current.0.clone();
166 Some(RefMulti::new(guard, k, v.get()))
167 };
168 }
169 }
170
171 if self.shard_i == self.map._shard_count() {
172 return None;
173 }
174
175 let guard = unsafe { self.map._yield_read_shard(self.shard_i) };
176
177 let iter = unsafe { guard.iter() };
178
179 self.current = Some((Arc::new(guard), iter));
180
181 self.shard_i += 1;
182 }
183 }
184}
185
186pub struct IterMut<'a, K, V, S = RandomState, M = DashMap<K, V, S>> {
199 map: &'a M,
200 shard_i: usize,
201 current: Option<GuardIterMut<'a, K, V>>,
202 marker: PhantomData<S>,
203}
204
205unsafe impl<'a, 'i, K, V, S, M> Send for IterMut<'i, K, V, S, M>
206where
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
214unsafe impl<'a, 'i, K, V, S, M> Sync for IterMut<'i, K, V, S, M>
215where
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
223impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>>
224 IterMut<'a, K, V, S, M>
225{
226 pub(crate) fn new(map: &'a M) -> Self {
227 Self {
228 map,
229 shard_i: 0,
230 current: None,
231 marker: PhantomData,
232 }
233 }
234}
235
236impl<'a, K: Eq + Hash, V, S: 'a + BuildHasher + Clone, M: Map<'a, K, V, S>> Iterator
237 for IterMut<'a, K, V, S, M>
238{
239 type Item = RefMutMulti<'a, K, V>;
240
241 fn next(&mut self) -> Option<Self::Item> {
242 loop {
243 if let Some(current) = self.current.as_mut() {
244 if let Some(b) = current.1.next() {
245 return unsafe {
246 let (k, v) = b.as_mut();
247 let guard = current.0.clone();
248 Some(RefMutMulti::new(guard, k, v.get_mut()))
249 };
250 }
251 }
252
253 if self.shard_i == self.map._shard_count() {
254 return None;
255 }
256
257 let guard = unsafe { self.map._yield_write_shard(self.shard_i) };
258
259 let iter = unsafe { guard.iter() };
260
261 self.current = Some((Arc::new(guard), iter));
262
263 self.shard_i += 1;
264 }
265 }
266}
267
268#[cfg(test)]
269mod tests {
270 use crate::DashMap;
271
272 #[test]
273 fn iter_mut_manual_count() {
274 let map = DashMap::new();
275
276 map.insert("Johnny", 21);
277
278 assert_eq!(map.len(), 1);
279
280 let mut c = 0;
281
282 for shard in map.shards() {
283 c += unsafe { shard.write().iter().count() };
284 }
285
286 assert_eq!(c, 1);
287 }
288
289 #[test]
290 fn iter_mut_count() {
291 let map = DashMap::new();
292
293 map.insert("Johnny", 21);
294
295 assert_eq!(map.len(), 1);
296
297 assert_eq!(map.iter_mut().count(), 1);
298 }
299
300 #[test]
301 fn iter_count() {
302 let map = DashMap::new();
303
304 map.insert("Johnny", 21);
305
306 assert_eq!(map.len(), 1);
307
308 assert_eq!(map.iter().count(), 1);
309 }
310}