Bevy PipeOrBreak

pipe_or_break.rs

pub struct PipeOrBreak<COut>(PhantomData<COut>);

impl<ACOut, A, B> Combine<A, B> for PipeOrBreak<ACOut>
where
    A: System<Out = ControlFlow<(), ACOut>>,
    B: System<In = ACOut>,
{
    type In = A::In;
    type Out = ();

    fn combine(
        input: Self::In,
        a: impl FnOnce(A::In) -> A::Out,
        b: impl FnOnce(B::In) -> B::Out,
    ) -> Self::Out {
        match a(input) {
            ControlFlow::Continue(out) => {
                b(out);
            }
            ControlFlow::Break(()) => (),
        }
    }
}

pub type PipeOrBreakSystem<ACOut, SystemA, SystemB> =
    CombinatorSystem<PipeOrBreak<ACOut>, SystemA, SystemB>;

pub trait PipeOrBreakExt<In, COut, Marker>:
    IntoSystem<In, ControlFlow<(), COut>, Marker>
{
    fn pipe_or_break<B, Final, MarkerB>(
        self,
        system_b: B,
    ) -> PipeOrBreakSystem<COut, Self::System, B::System>
    where
        B: IntoSystem<COut, Final, MarkerB>;
}

impl<T, In, COut, Marker> PipeOrBreakExt<In, COut, Marker> for T
where
    T: IntoSystem<In, ControlFlow<(), COut>, Marker>,
{
    fn pipe_or_break<B, Final, MarkerB>(
        self,
        system_b: B,
    ) -> PipeOrBreakSystem<COut, Self::System, B::System>
    where
        B: IntoSystem<COut, Final, MarkerB>,
    {
        let system_a = IntoSystem::into_system(self);
        let system_b = IntoSystem::into_system(system_b);
        let name =
            format!("PipeOrBreak({}, {})", system_a.name(), system_b.name());
        PipeOrBreakSystem::new(system_a, system_b, Cow::Owned(name))
    }
}

pub fn break_if_none<T>(a: Option<T>) -> ControlFlow<(), T> {
    match a {
        Some(a) => ControlFlow::Continue(a),
        None => ControlFlow::Break(()),
    }
}

// === [ Usage ] =====================================================================
fn main() {
    App::new()
        .add_systems(
            Update,
            a_system_that_returns_option
                .map(break_if_none)
                .pipe_or_break(another_system)
        )
}