Skip to content

feat: Add parallel iterators #10075

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
Mar 25, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 22 additions & 2 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

19 changes: 19 additions & 0 deletions crates/swc_par_iter/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
[package]
description = "Fork of rayon, with chili support"
edition = { workspace = true }
license = { workspace = true }
name = "swc_par_iter"
repository = { workspace = true }
version = "1.0.0"

[dependencies]
either = { workspace = true }
swc_parallel = { version = "1.3.0", path = "../swc_parallel", default-features = false }

[dev-dependencies]
chili = { workspace = true }
rand = { workspace = true }
rand_xorshift = "0.3"
swc_parallel = { version = "1.3.0", path = "../swc_parallel", default-features = false, features = [
"chili",
] }
87 changes: 87 additions & 0 deletions crates/swc_par_iter/src/array.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
//! Parallel iterator types for [arrays] (`[T; N]`)
//!
//! You will rarely need to interact with this module directly unless you need
//! to name one of the iterator types.
//!
//! [arrays]: https://doc.rust-lang.org/std/primitive.array.html

use std::mem::ManuallyDrop;

use crate::{
iter::{plumbing::*, *},
slice::{Iter, IterMut},
vec::DrainProducer,
};

impl<'data, T: Sync + 'data, const N: usize> IntoParallelIterator for &'data [T; N] {
type Item = &'data T;
type Iter = Iter<'data, T>;

fn into_par_iter(self) -> Self::Iter {
<&[T]>::into_par_iter(self)
}
}

impl<'data, T: Send + 'data, const N: usize> IntoParallelIterator for &'data mut [T; N] {
type Item = &'data mut T;
type Iter = IterMut<'data, T>;

fn into_par_iter(self) -> Self::Iter {
<&mut [T]>::into_par_iter(self)
}
}

impl<T: Send, const N: usize> IntoParallelIterator for [T; N] {
type Item = T;
type Iter = IntoIter<T, N>;

fn into_par_iter(self) -> Self::Iter {
IntoIter { array: self }
}
}

/// Parallel iterator that moves out of an array.
#[derive(Debug, Clone)]
pub struct IntoIter<T: Send, const N: usize> {
array: [T; N],
}

impl<T: Send, const N: usize> ParallelIterator for IntoIter<T, N> {
type Item = T;

fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
bridge(self, consumer)
}

fn opt_len(&self) -> Option<usize> {
Some(N)
}
}

impl<T: Send, const N: usize> IndexedParallelIterator for IntoIter<T, N> {
fn drive<C>(self, consumer: C) -> C::Result
where
C: Consumer<Self::Item>,
{
bridge(self, consumer)
}

fn len(&self) -> usize {
N
}

fn with_producer<CB>(self, callback: CB) -> CB::Output
where
CB: ProducerCallback<Self::Item>,
{
unsafe {
// Drain every item, and then the local array can just fall out of scope.
let mut array = ManuallyDrop::new(self.array);
let producer = DrainProducer::new(array.as_mut_slice());
callback.callback(producer)
}
}
}
122 changes: 122 additions & 0 deletions crates/swc_par_iter/src/collections/binary_heap.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
//! This module contains the parallel iterator types for heaps
//! (`BinaryHeap<T>`). You will rarely need to interact with it directly
//! unless you have need to name one of the iterator types.

use std::collections::BinaryHeap;

use crate::{
iter::{plumbing::*, *},
vec,
};

/// Parallel iterator over a binary heap
#[derive(Debug, Clone)]
pub struct IntoIter<T: Ord + Send> {
inner: vec::IntoIter<T>,
}

impl<T: Ord + Send> IntoParallelIterator for BinaryHeap<T> {
type Item = T;
type Iter = IntoIter<T>;

fn into_par_iter(self) -> Self::Iter {
IntoIter {
inner: Vec::from(self).into_par_iter(),
}
}
}

delegate_indexed_iterator! {
IntoIter<T> => T,
impl<T: Ord + Send>
}

/// Parallel iterator over an immutable reference to a binary heap
#[derive(Debug)]
pub struct Iter<'a, T: Ord + Sync> {
// TODO (MSRV 1.80): this could use `slice::Iter` built from
// `BinaryHeap::as_slice`, rather than using a temp `Vec<&T>`.
inner: vec::IntoIter<&'a T>,
}

impl<'a, T: Ord + Sync> Clone for Iter<'a, T> {
fn clone(&self) -> Self {
Iter {
inner: self.inner.clone(),
}
}
}

into_par_vec! {
&'a BinaryHeap<T> => Iter<'a, T>,
impl<'a, T: Ord + Sync>
}

delegate_indexed_iterator! {
Iter<'a, T> => &'a T,
impl<'a, T: Ord + Sync + 'a>
}

// `BinaryHeap` doesn't have a mutable `Iterator`

/// Draining parallel iterator that moves out of a binary heap,
/// but keeps the total capacity.
#[derive(Debug)]
pub struct Drain<'a, T: Ord + Send> {
heap: &'a mut BinaryHeap<T>,
}

impl<'a, T: Ord + Send> ParallelDrainFull for &'a mut BinaryHeap<T> {
type Item = T;
type Iter = Drain<'a, T>;

fn par_drain(self) -> Self::Iter {
Drain { heap: self }
}
}

impl<'a, T: Ord + Send> ParallelIterator for Drain<'a, T> {
type Item = T;

fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
bridge(self, consumer)
}

fn opt_len(&self) -> Option<usize> {
Some(self.len())
}
}

impl<'a, T: Ord + Send> IndexedParallelIterator for Drain<'a, T> {
fn drive<C>(self, consumer: C) -> C::Result
where
C: Consumer<Self::Item>,
{
bridge(self, consumer)
}

fn len(&self) -> usize {
self.heap.len()
}

fn with_producer<CB>(self, callback: CB) -> CB::Output
where
CB: ProducerCallback<Self::Item>,
{
super::DrainGuard::new(self.heap)
.par_drain(..)
.with_producer(callback)
}
}

impl<'a, T: Ord + Send> Drop for Drain<'a, T> {
fn drop(&mut self) {
if !self.heap.is_empty() {
// We must not have produced, so just call a normal drain to remove the items.
self.heap.drain();
}
}
}
66 changes: 66 additions & 0 deletions crates/swc_par_iter/src/collections/btree_map.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
//! This module contains the parallel iterator types for B-Tree maps
//! (`BTreeMap<K, V>`). You will rarely need to interact with it directly
//! unless you have need to name one of the iterator types.

use std::collections::BTreeMap;

use crate::{
iter::{plumbing::*, *},
vec,
};

/// Parallel iterator over a B-Tree map
#[derive(Debug)] // std doesn't Clone
pub struct IntoIter<K: Ord + Send, V: Send> {
inner: vec::IntoIter<(K, V)>,
}

into_par_vec! {
BTreeMap<K, V> => IntoIter<K, V>,
impl<K: Ord + Send, V: Send>
}

delegate_iterator! {
IntoIter<K, V> => (K, V),
impl<K: Ord + Send, V: Send>
}

/// Parallel iterator over an immutable reference to a B-Tree map
#[derive(Debug)]
pub struct Iter<'a, K: Ord + Sync, V: Sync> {
inner: vec::IntoIter<(&'a K, &'a V)>,
}

impl<'a, K: Ord + Sync, V: Sync> Clone for Iter<'a, K, V> {
fn clone(&self) -> Self {
Iter {
inner: self.inner.clone(),
}
}
}

into_par_vec! {
&'a BTreeMap<K, V> => Iter<'a, K, V>,
impl<'a, K: Ord + Sync, V: Sync>
}

delegate_iterator! {
Iter<'a, K, V> => (&'a K, &'a V),
impl<'a, K: Ord + Sync + 'a, V: Sync + 'a>
}

/// Parallel iterator over a mutable reference to a B-Tree map
#[derive(Debug)]
pub struct IterMut<'a, K: Ord + Sync, V: Send> {
inner: vec::IntoIter<(&'a K, &'a mut V)>,
}

into_par_vec! {
&'a mut BTreeMap<K, V> => IterMut<'a, K, V>,
impl<'a, K: Ord + Sync, V: Send>
}

delegate_iterator! {
IterMut<'a, K, V> => (&'a K, &'a mut V),
impl<'a, K: Ord + Sync + 'a, V: Send + 'a>
}
Loading
Loading