Skip to content

Commit

Permalink
Merge pull request indexmap-rs#197 from cuviper/par_drain
Browse files Browse the repository at this point in the history
impl ParallelDrainRange
  • Loading branch information
cuviper authored Jan 6, 2022
2 parents f9a3ca9 + 1426285 commit 4d6dde3
Show file tree
Hide file tree
Showing 6 changed files with 93 additions and 5 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ bench = false
autocfg = "1"
[dependencies]
serde = { version = "1.0", optional = true, default-features = false }
rayon = { version = "1.2", optional = true }
rayon = { version = "1.4.1", optional = true }

# Internal feature, only used when building as part of rustc,
# not part of the stable interface of this crate.
Expand Down
4 changes: 2 additions & 2 deletions src/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,12 +69,12 @@ pub use self::core::{Entry, OccupiedEntry, VacantEntry};
/// ```
#[cfg(has_std)]
pub struct IndexMap<K, V, S = RandomState> {
core: IndexMapCore<K, V>,
pub(crate) core: IndexMapCore<K, V>,
hash_builder: S,
}
#[cfg(not(has_std))]
pub struct IndexMap<K, V, S> {
core: IndexMapCore<K, V>,
pub(crate) core: IndexMapCore<K, V>,
hash_builder: S,
}

Expand Down
13 changes: 13 additions & 0 deletions src/map/core.rs
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,19 @@ impl<K, V> IndexMapCore<K, V> {
self.entries.drain(range)
}

#[cfg(feature = "rayon")]
pub(crate) fn par_drain<R>(&mut self, range: R) -> rayon::vec::Drain<'_, Bucket<K, V>>
where
K: Send,
V: Send,
R: RangeBounds<usize>,
{
use rayon::iter::ParallelDrainRange;
let range = simplify_range(range, self.entries.len());
self.erase_indices(range.start, range.end);
self.entries.par_drain(range)
}

pub(crate) fn split_off(&mut self, at: usize) -> Self {
assert!(at <= self.entries.len());
self.erase_indices(at, self.entries.len());
Expand Down
38 changes: 38 additions & 0 deletions src/rayon/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use crate::vec::Vec;
use core::cmp::Ordering;
use core::fmt;
use core::hash::{BuildHasher, Hash};
use core::ops::RangeBounds;

use crate::Bucket;
use crate::Entries;
Expand Down Expand Up @@ -156,6 +157,43 @@ impl<K: Sync + Send, V: Send> IndexedParallelIterator for ParIterMut<'_, K, V> {
indexed_parallel_iterator_methods!(Bucket::ref_mut);
}

/// Requires crate feature `"rayon"`.
impl<'a, K, V, S> ParallelDrainRange<usize> for &'a mut IndexMap<K, V, S>
where
K: Send,
V: Send,
{
type Item = (K, V);
type Iter = ParDrain<'a, K, V>;

fn par_drain<R: RangeBounds<usize>>(self, range: R) -> Self::Iter {
ParDrain {
entries: self.core.par_drain(range),
}
}
}

/// A parallel draining iterator over the entries of a `IndexMap`.
///
/// This `struct` is created by the [`par_drain`] method on [`IndexMap`]
/// (provided by rayon's `ParallelDrainRange` trait). See its documentation for more.
///
/// [`par_drain`]: ../struct.IndexMap.html#method.par_drain
/// [`IndexMap`]: ../struct.IndexMap.html
pub struct ParDrain<'a, K: Send, V: Send> {
entries: rayon::vec::Drain<'a, Bucket<K, V>>,
}

impl<K: Send, V: Send> ParallelIterator for ParDrain<'_, K, V> {
type Item = (K, V);

parallel_iterator_methods!(Bucket::key_value);
}

impl<K: Send, V: Send> IndexedParallelIterator for ParDrain<'_, K, V> {
indexed_parallel_iterator_methods!(Bucket::key_value);
}

/// Parallel iterator methods and other parallel methods.
///
/// The following methods **require crate feature `"rayon"`**.
Expand Down
37 changes: 37 additions & 0 deletions src/rayon/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use crate::vec::Vec;
use core::cmp::Ordering;
use core::fmt;
use core::hash::{BuildHasher, Hash};
use core::ops::RangeBounds;

use crate::Entries;
use crate::IndexSet;
Expand Down Expand Up @@ -111,6 +112,42 @@ impl<T: Sync> IndexedParallelIterator for ParIter<'_, T> {
indexed_parallel_iterator_methods!(Bucket::key_ref);
}

/// Requires crate feature `"rayon"`.
impl<'a, T, S> ParallelDrainRange<usize> for &'a mut IndexSet<T, S>
where
T: Send,
{
type Item = T;
type Iter = ParDrain<'a, T>;

fn par_drain<R: RangeBounds<usize>>(self, range: R) -> Self::Iter {
ParDrain {
entries: self.map.core.par_drain(range),
}
}
}

/// A parallel draining iterator over the items of a `IndexSet`.
///
/// This `struct` is created by the [`par_drain`] method on [`IndexSet`]
/// (provided by rayon's `ParallelDrainRange` trait). See its documentation for more.
///
/// [`par_drain`]: ../struct.IndexSet.html#method.par_drain
/// [`IndexSet`]: ../struct.IndexSet.html
pub struct ParDrain<'a, T: Send> {
entries: rayon::vec::Drain<'a, Bucket<T>>,
}

impl<T: Send> ParallelIterator for ParDrain<'_, T> {
type Item = T;

parallel_iterator_methods!(Bucket::key);
}

impl<T: Send> IndexedParallelIterator for ParDrain<'_, T> {
indexed_parallel_iterator_methods!(Bucket::key);
}

/// Parallel iterator methods and other parallel methods.
///
/// The following methods **require crate feature `"rayon"`**.
Expand Down
4 changes: 2 additions & 2 deletions src/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,11 +61,11 @@ type Bucket<T> = super::Bucket<T, ()>;
/// ```
#[cfg(has_std)]
pub struct IndexSet<T, S = RandomState> {
map: IndexMap<T, (), S>,
pub(crate) map: IndexMap<T, (), S>,
}
#[cfg(not(has_std))]
pub struct IndexSet<T, S> {
map: IndexMap<T, (), S>,
pub(crate) map: IndexMap<T, (), S>,
}

impl<T, S> Clone for IndexSet<T, S>
Expand Down

0 comments on commit 4d6dde3

Please sign in to comment.