Add interface to rayon's par_bridge adapter

SpectralSequences · Apr 5, 2024 · 03b69cd · 03b69cd
1 parent 1c09d2a
commit 03b69cd
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/ext/crates/maybe-rayon/src/concurrent.rs b/ext/crates/maybe-rayon/src/concurrent.rs
@@ -14,6 +14,14 @@ pub mod prelude {
         fn maybe_par_iter_mut(&'data mut self) -> Self::Iter;
     }
 
+    pub type MaybeIterBridge<I> = rayon::iter::IterBridge<I>;
+
+    pub trait MaybeParallelBridge: ParallelBridge {
+        fn maybe_par_bridge(self) -> MaybeIterBridge<Self> {
+            self.par_bridge()
+        }
+    }
+
     // Implementations
 
     impl<I: ParallelIterator> MaybeParallelIterator for I {}
@@ -33,6 +41,8 @@ pub mod prelude {
             self.par_iter_mut()
         }
     }
+
+    impl<I: ParallelBridge> MaybeParallelBridge for I {}
 }
 
 pub fn join<A, B, RA, RB>(oper_a: A, oper_b: B) -> (RA, RB)

diff --git a/ext/crates/maybe-rayon/src/sequential.rs b/ext/crates/maybe-rayon/src/sequential.rs
@@ -15,6 +15,14 @@ pub mod prelude {
         fn maybe_par_iter_mut(&'data mut self) -> Self::Iter;
     }
 
+    pub struct MaybeIterBridge<Iter>(Iter);
+
+    pub trait MaybeParallelBridge: Sized {
+        fn maybe_par_bridge(self) -> MaybeIterBridge<Self> {
+            MaybeIterBridge(self)
+        }
+    }
+
     // Implementations
 
     impl<I: Iterator> MaybeParallelIterator for I {}
@@ -39,6 +47,16 @@ pub mod prelude {
             self.into_iter()
         }
     }
+
+    impl<Iter: Iterator> Iterator for MaybeIterBridge<Iter> {
+        type Item = Iter::Item;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            self.0.next()
+        }
+    }
+
+    impl<T: Iterator> MaybeParallelBridge for T {}
 }
 
 #[allow(dead_code)]