use std::marker::PhantomData;
use zrx_scheduler::action::context::Binding;
use zrx_scheduler::action::{Action, Context};
use zrx_scheduler::step::IntoSteps;
use zrx_scheduler::{Id, Value};
use crate::stream::function::{Arguments, MapFn};
use crate::stream::Stream;
use super::Operator;
pub struct Map<T, F, A, U> {
function: F,
marker: PhantomData<(T, A, U)>,
}
impl<I, T> Stream<I, T>
where
I: Id,
T: Value,
{
#[inline]
pub fn map<F, A, U>(&self, f: F) -> Stream<I, U>
where
F: MapFn<A, I, T, U> + Clone,
A: Arguments,
U: Value,
{
self.subscribe(Map {
function: f,
marker: PhantomData,
})
}
}
impl<I, T, F, A, U> Action<I> for Map<T, F, A, U>
where
I: Id,
T: Value,
F: MapFn<A, I, T, U> + Clone,
A: Arguments,
U: Value,
{
type Inputs = (T,);
type Output<'a> = U;
fn execute(&mut self, ctx: Context<I, Self>) -> impl IntoSteps<I, Self> {
let Binding { scopes, inputs, mut output, .. } = ctx.bind();
scopes.into_iter().map(move |mut scope| {
let Some(value) = inputs.get(scope.key()).cloned() else {
output.remove(scope.key());
return scope.done();
};
scope.task().build({
let function = self.function.clone();
move || {
let value = function.execute(&mut scope, value)?;
scope.then().build(move |mut ctx| {
let mut output = ctx.output().expect("invariant");
output.insert(scope.key().clone(), value);
scope.done()
})
}
})
})
}
}