diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 5e38fed..71fd4ad 100644 --- a/.github/workflows/rust.yml +++ b/.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 diff --git a/src/iter.rs b/src/iter.rs new file mode 100644 index 0000000..70d2ef1 --- /dev/null +++ b/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, + 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) -> 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 { + 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::>(); + 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::::new(); + + let mut iter = map.iter(); + assert_eq!(iter.next(), None); + } +} diff --git a/src/lib.rs b/src/lib.rs index cb39fde..eaab1e1 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/shard.rs b/src/shard.rs index 732ebff..b71e7ea 100644 --- a/src/shard.rs +++ b/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 { + self.data.write().unwrap() + } + + pub fn read_sync(&self) -> ShardReader { + self.data.read().unwrap() + } } impl std::ops::Deref for Shard { diff --git a/src/shard_map.rs b/src/shard_map.rs index bfe854d..e14fd33 100644 --- a/src/shard_map.rs +++ b/src/shard_map.rs @@ -185,7 +185,7 @@ where } #[inline] - fn shard(&self, key: &K) -> (&CachePadded>, u64) { + pub(crate) fn shard(&self, key: &K) -> (&CachePadded>, u64) { let hash = self.hash_u64(key); let shard_idx = self.shard_for_hash(hash as usize); @@ -193,6 +193,20 @@ where (unsafe { self.inner.shards.get_unchecked(shard_idx) }, hash) } + #[inline] + pub(crate) fn shard_by_idx(&self, idx: usize) -> &CachePadded> { + 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 { + 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. ///