feat: add iterators

.github/workflows/rust.yml

@@ -12,12 +12,6 @@ on:
 
   pull_request:
     branches: [ "main" ]
-    paths:
-      - src/
-      - tests/
-      - Cargo.toml
-      - Cargo.lock
-      - .github/
 
 env:
   CARGO_TERM_COLOR: always

src/iter.rs

@@ -0,0 +1,103 @@
+use std::hash::BuildHasher;
+
+enum IterState<'a, K, V> {
+    Start,
+    Shard(
+        usize,
+        hashbrown::raw::RawIter<(K, V)>,
+        crate::shard::ShardReader<'a, K, V>,
+    ),
+    Finished,
+}
+
+pub struct Iter<'a, K, V, S> {
+    map: crate::ShardMap<K, V, S>,
+    state: IterState<'a, K, V>,
+}
+
+impl<'a, K: 'static, V: 'static, S: BuildHasher> Iter<'a, K, V, S>
+where
+    K: Eq + std::hash::Hash + 'static,
+    V: 'static,
+{
+    pub fn new(map: crate::ShardMap<K, V, S>) -> Self {
+        Self {
+            map,
+            state: IterState::Start,
+        }
+    }
+}
+
+impl<'a, K: Clone + 'static, V: 'static, S: BuildHasher> Iterator for Iter<'a, K, V, S>
+where
+    K: Eq + std::hash::Hash + 'static,
+    V: 'static,
+{
+    type Item = (&'a K, &'a V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match &mut self.state {
+            IterState::Start => {
+                let map = self.map.shard_by_idx(0);
+
+                let reader: std::sync::RwLockReadGuard<'static, hashbrown::raw::RawTable<(K, V)>> =
+                    unsafe { std::mem::transmute(map.read_sync()) };
+                self.state = IterState::Shard(0, unsafe { reader.iter() }, reader);
+
+                self.next()
+            }
+            IterState::Shard(idx, iter, _shard) => {
+                match iter
+                    .next()
+                    .map(|bucket| unsafe { bucket.as_ref() })
+                    .map(|bucket| (&bucket.0, &bucket.1))
+                {
+                    Some(value) => Some(value),
+                    None => {
+                        if *idx >= self.map.num_shards() - 1 {
+                            self.state = IterState::Finished;
+                            return None;
+                        }
+
+                        let map = self.map.shard_by_idx(*idx + 1);
+
+                        let reader: std::sync::RwLockReadGuard<
+                            'static,
+                            hashbrown::raw::RawTable<(K, V)>,
+                        > = unsafe { std::mem::transmute(map.read_sync()) };
+
+                        self.state = IterState::Shard(*idx + 1, unsafe { reader.iter() }, reader);
+
+                        self.next()
+                    }
+                }
+            }
+            IterState::Finished => None,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[tokio::test]
+    async fn test_iter() {
+        let map = crate::ShardMap::new();
+        map.insert("foo", "bar").await;
+        map.insert("baz", "qux").await;
+        let mut iter = map.iter().collect::<Vec<_>>();
+        iter.sort_by(|a, b| a.0.cmp(b.0));
+
+        assert_eq!(iter[0], (&"baz", &"qux"));
+        assert_eq!(iter[1], (&"foo", &"bar"));
+
+        assert_eq!(iter.len(), 2);
+    }
+
+    #[tokio::test]
+    async fn test_iter_empty() {
+        let map = crate::ShardMap::<u32, u32>::new();
+
+        let mut iter = map.iter();
+        assert_eq!(iter.next(), None);
+    }
+}

src/lib.rs

@@ -53,6 +53,7 @@
 //!
 //! See the documentation for each data structure for more information.
 
+pub mod iter;
 pub mod mapref;
 mod shard;
 mod shard_map;

src/shard.rs

@@ -30,6 +30,14 @@ where
     pub async fn read<'a>(&'a self) -> ShardReader<'a, K, V> {
         Read::new(self).await
     }
+
+    pub fn write_sync(&self) -> ShardWriter<K, V> {
+        self.data.write().unwrap()
+    }
+
+    pub fn read_sync(&self) -> ShardReader<K, V> {
+        self.data.read().unwrap()
+    }
 }
 
 impl<K, V> std::ops::Deref for Shard<K, V> {

src/shard_map.rs

@@ -185,14 +185,28 @@ where
     }
 
     #[inline]
-    fn shard(&self, key: &K) -> (&CachePadded<Shard<K, V>>, u64) {
+    pub(crate) fn shard(&self, key: &K) -> (&CachePadded<Shard<K, V>>, u64) {
         let hash = self.hash_u64(key);
 
         let shard_idx = self.shard_for_hash(hash as usize);
 
         (unsafe { self.inner.shards.get_unchecked(shard_idx) }, hash)
     }
 
+    #[inline]
+    pub(crate) fn shard_by_idx(&self, idx: usize) -> &CachePadded<Shard<K, V>> {
+        unsafe { self.inner.shards.get_unchecked(idx) }
+    }
+
+    #[inline]
+    pub(crate) fn num_shards(&self) -> usize {
+        self.inner.len()
+    }
+
+    pub fn iter(&self) -> crate::iter::Iter<K, V, S> {
+        crate::iter::Iter::new(self.clone())
+    }
+
     /// Inserts a key-value pair into the map. If the key already exists, the value is updated and
     /// the old value is returned.
     ///