use crate::pipe::fan_in::{Exhausted, FanIn, FanInStrategy};
use crate::pipe::primitives::{AndThen, Pipe, UnpinPipe};
use proxima_core::markers::DropSafe;
#[cfg(feature = "alloc")]
use crate::pipe::fanout::{AllOrNothing, FanOut};
#[cfg(feature = "alloc")]
use alloc::vec;
pub trait PipeExt: Pipe + Sized {
fn and_then<Next>(self, next: Next) -> AndThen<Self, Next>
where
Next: Pipe<In = Self::Out>,
Next::Err: From<Self::Err>,
{
AndThen::new(self, next)
}
fn filter<Pred>(self, predicate: Pred) -> AndThen<Pred, Self>
where
Pred: Pipe<Out = Self::In>,
Self::Err: From<Pred::Err>,
{
AndThen::new(predicate, self)
}
#[cfg(feature = "alloc")]
fn fanout(self, other: Self) -> FanOut<Self, AllOrNothing>
where
Self: Clone,
Self::In: Clone,
{
FanOut::new(vec![self, other])
}
fn fanin<Strategy>(self, other: Self, strategy: Strategy) -> FanIn<Self, Strategy, 2>
where
Self: UnpinPipe<In = (), Err = Exhausted> + DropSafe,
Strategy: FanInStrategy,
{
FanIn::new([self, other], strategy)
}
}
impl<P: Pipe> PipeExt for P {}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::{FanInStrategy, Pipe, PipeExt};
use crate::block_on;
use core::future::Future;
#[derive(Debug, PartialEq, Eq)]
struct Overflow;
struct Increment;
impl Pipe for Increment {
type In = u64;
type Out = u64;
type Err = Overflow;
fn call(&self, input: u64) -> impl Future<Output = Result<u64, Overflow>> {
async move { input.checked_add(1).ok_or(Overflow) }
}
}
#[derive(Debug, PartialEq, Eq)]
enum ChainError {
Overflowed,
Rejected,
}
impl From<Overflow> for ChainError {
fn from(_: Overflow) -> Self {
ChainError::Overflowed
}
}
struct Double;
impl Pipe for Double {
type In = u64;
type Out = u64;
type Err = ChainError;
fn call(&self, input: u64) -> impl Future<Output = Result<u64, ChainError>> {
async move { Ok(input * 2) }
}
}
struct RejectOdd;
impl Pipe for RejectOdd {
type In = u64;
type Out = u64;
type Err = ChainError;
fn call(&self, input: u64) -> impl Future<Output = Result<u64, ChainError>> {
async move {
if input.is_multiple_of(2) {
Ok(input)
} else {
Err(ChainError::Rejected)
}
}
}
}
#[test]
fn and_then_and_filter_compose_through_the_ext_sugar() {
let pipeline = Increment.and_then(Double).filter(RejectOdd);
let admitted = block_on(Pipe::call(&pipeline, 4));
assert_eq!(admitted, Ok(10), "4 -> RejectOdd passes -> +1 -> *2 -> 10");
let rejected = block_on(Pipe::call(&pipeline, 3));
assert_eq!(
rejected,
Err(ChainError::Rejected),
"3 is odd: RejectOdd stops the chain before Increment runs"
);
}
#[cfg(feature = "alloc")]
mod fanout_ext {
use super::{Pipe, PipeExt};
use alloc::sync::Arc;
use core::convert::Infallible;
use core::future::Future;
use core::sync::atomic::{AtomicUsize, Ordering};
#[derive(Clone)]
struct RecordingSink {
calls: Arc<AtomicUsize>,
}
impl Pipe for RecordingSink {
type In = u32;
type Out = ();
type Err = Infallible;
fn call(&self, _input: u32) -> impl Future<Output = Result<(), Infallible>> {
self.calls.fetch_add(1, Ordering::Relaxed);
async move { Ok(()) }
}
}
#[test]
fn fanout_ext_broadcasts_to_both_arms() {
let calls = Arc::new(AtomicUsize::new(0));
let fan = RecordingSink {
calls: Arc::clone(&calls),
}
.fanout(RecordingSink {
calls: Arc::clone(&calls),
});
let outcome = super::block_on(Pipe::call(&fan, 7));
assert_eq!(outcome, Ok(()));
assert_eq!(calls.load(Ordering::Relaxed), 2, "both arms ran");
}
}
struct OnceSource {
emitted: core::cell::Cell<bool>,
value: u32,
}
impl proxima_core::markers::DropSafe for OnceSource {}
impl super::UnpinPipe for OnceSource {
type In = ();
type Out = u32;
type Err = super::Exhausted;
fn call(&self, (): ()) -> impl Future<Output = Result<u32, super::Exhausted>> + Unpin {
if self.emitted.replace(true) {
core::future::ready(Err(super::Exhausted))
} else {
core::future::ready(Ok(self.value))
}
}
}
impl Pipe for OnceSource {
type In = ();
type Out = u32;
type Err = super::Exhausted;
fn call(&self, (): ()) -> impl Future<Output = Result<u32, super::Exhausted>> {
super::UnpinPipe::call(self, ())
}
}
#[test]
fn fanin_ext_merges_both_sources_via_the_sugar() {
let fan = (OnceSource {
emitted: core::cell::Cell::new(false),
value: 1,
})
.fanin(
OnceSource {
emitted: core::cell::Cell::new(false),
value: 2,
},
super::super::fan_in::Select::RoundRobin,
);
let waker = core::task::Waker::noop();
let mut context = core::task::Context::from_waker(waker);
let mut merged = alloc_free_collect(&fan, &mut context);
merged.sort_unstable();
assert_eq!(merged, [1, 2], "both sources drain through the fanin sugar");
}
fn alloc_free_collect<S, Strategy, const N: usize>(
fan: &super::FanIn<S, Strategy, N>,
cx: &mut core::task::Context<'_>,
) -> [u32; N]
where
S: super::UnpinPipe<In = (), Out = u32, Err = super::Exhausted>
+ proxima_core::markers::DropSafe,
Strategy: FanInStrategy,
{
let mut out = [0u32; N];
let mut count = 0;
for _ in 0..(N * 4) {
if count == N {
break;
}
let mut call = Pipe::call(fan, ());
match core::pin::Pin::new(&mut call).poll(cx) {
core::task::Poll::Ready(Ok(value)) => {
out[count] = value;
count += 1;
}
core::task::Poll::Ready(Err(super::Exhausted)) => break,
core::task::Poll::Pending => {}
}
}
out
}
}