use async_trait::async_trait;
use futures::channel::mpsc;
use futures::stream::{FuturesUnordered, StreamExt};
use pipecrab_core::{DataFrame, Decision, Direction, Processor, SystemFrame};
use crate::{Inbound, MaybeSend, Outbound, Stage, StageError};
#[cfg(not(target_arch = "wasm32"))]
pub type DriverFuture = futures::future::BoxFuture<'static, ()>;
#[cfg(target_arch = "wasm32")]
pub type DriverFuture = futures::future::LocalBoxFuture<'static, ()>;
const DEFAULT_CAPACITY: usize = 16;
pub fn link(capacity: usize) -> (Outbound, Inbound) {
let capacity = capacity.max(1);
let (data_tx, data_rx) = mpsc::channel::<DataFrame>(capacity);
let (sys_tx, sys_rx) = mpsc::channel::<(Direction, SystemFrame)>(capacity);
(Outbound { data: data_tx, sys: sys_tx }, Inbound { sys: sys_rx, data: data_rx })
}
pub struct PipelineEnds {
pub input: Outbound,
pub output: Inbound,
}
pub struct PipelineBuilder<E> {
stages: Vec<Box<dyn Stage<Effect = E>>>,
capacity: usize,
}
impl<E: MaybeSend + 'static> PipelineBuilder<E> {
pub fn new() -> Self {
Self { stages: Vec::new(), capacity: DEFAULT_CAPACITY }
}
pub fn capacity(mut self, capacity: usize) -> Self {
self.capacity = capacity.max(1);
self
}
pub fn stage(mut self, stage: impl Stage<Effect = E> + 'static) -> Self {
self.stages.push(Box::new(stage));
self
}
pub fn boxed(mut self, stage: Box<dyn Stage<Effect = E>>) -> Self {
self.stages.push(stage);
self
}
pub fn build(self) -> Pipeline<E> {
assert!(!self.stages.is_empty(), "a pipeline needs at least one stage");
Pipeline { stages: self.stages, capacity: self.capacity }
}
}
impl<E: MaybeSend + 'static> Default for PipelineBuilder<E> {
fn default() -> Self {
Self::new()
}
}
pub struct Pipeline<E> {
stages: Vec<Box<dyn Stage<Effect = E>>>,
capacity: usize,
}
impl<E: MaybeSend + 'static> Pipeline<E> {
pub fn start(self) -> (PipelineEnds, DriverFuture) {
let capacity = self.capacity;
let (input, head_in) = link(capacity);
let (tail_out, output) = link(capacity);
let driver = Box::new(self).run(head_in, tail_out);
(PipelineEnds { input, output }, driver)
}
}
impl<E> Processor for Pipeline<E> {
type Effect = E;
fn decide_data(&mut self, _frame: &DataFrame) -> Decision<E> {
unreachable!("a Pipeline is driven by Stage::run, not decide_data")
}
fn decide_system(&mut self, _dir: Direction, _frame: &SystemFrame) -> Decision<E> {
unreachable!("a Pipeline is driven by Stage::run, not decide_system")
}
}
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl<E: MaybeSend + 'static> Stage for Pipeline<E> {
async fn perform(&self, _effect: E, _out: &Outbound) -> Result<(), StageError> {
unreachable!("a Pipeline is driven by Stage::run, not perform")
}
async fn run(self: Box<Self>, inbound: Inbound, out: Outbound) {
let Pipeline { stages, capacity } = *self;
let n = stages.len();
let mut tasks = FuturesUnordered::new();
let mut current_in = Some(inbound);
let mut final_out = Some(out);
for (i, stage) in stages.into_iter().enumerate() {
let stage_in = current_in.take().expect("inbound threaded through");
let stage_out = if i + 1 == n {
final_out.take().expect("outbound threaded through")
} else {
let (this_out, next_in) = link(capacity);
current_in = Some(next_in);
this_out
};
tasks.push(stage.run(stage_in, stage_out));
}
while tasks.next().await.is_some() {}
}
}