[Merged by Bors] - Add AND/OR combinators for run conditions

crates/bevy_ecs/src/schedule/condition.rs

@@ -1,12 +1,41 @@
-use crate::system::BoxedSystem;
+use std::borrow::Cow;
+
+use crate::{
+    system::{BoxedSystem, CombinatorSystem, Combine, IntoSystem, System},
+    world::World,
+};
 
 pub type BoxedCondition = BoxedSystem<(), bool>;
 
 /// A system that determines if one or more scheduled systems should run.
 ///
 /// Implemented for functions and closures that convert into [`System<In=(), Out=bool>`](crate::system::System)
 /// with [read-only](crate::system::ReadOnlySystemParam) parameters.
-pub trait Condition<Params>: sealed::Condition<Params> {}
+pub trait Condition<Params>: sealed::Condition<Params> {
+    /// Returns a new run condition that only returns `true`
+    /// if both this one and the passed `and_then` return `true`.
+    ///
+    /// The returned run condition is short-circuiting, meaning
+    /// `and_then` will only be invoked if `self` returns `true`.
+    fn and_then<P, C: Condition<P>>(self, and_then: C) -> AndThen<Self::System, C::System> {
+        let a = IntoSystem::into_system(self);
+        let b = IntoSystem::into_system(and_then);
+        let name = format!("{} && {}", a.name(), b.name());
+        CombinatorSystem::new(a, b, Cow::Owned(name))
+    }
+
+    /// Returns a new run condition that returns `true`
+    /// if either this one or the passed `or_else` return `true`.
+    ///
+    /// The returned run condition is short-circuiting, meaning
+    /// `and_then` will only be invoked if `self` returns `false`.
+    fn or_else<P, C: Condition<P>>(self, or_else: C) -> OrElse<Self::System, C::System> {
+        let a = IntoSystem::into_system(self);
+        let b = IntoSystem::into_system(or_else);
+        let name = format!("{} || {}", a.name(), b.name());
+        CombinatorSystem::new(a, b, Cow::Owned(name))
+    }
+}
 
 impl<Params, F> Condition<Params> for F where F: sealed::Condition<Params> {}
 
@@ -142,3 +171,69 @@ pub mod common_conditions {
         condition.pipe(|In(val): In<bool>| !val)
     }
 }
+
+pub type AndThen<A, B> = CombinatorSystem<AndThenMarker, A, B>;
+
+pub type OrElse<A, B> = CombinatorSystem<OrElseMarker, A, B>;
+
+#[doc(hidden)]
+pub struct AndThenMarker;
+
+impl<In, A, B> Combine<A, B> for AndThenMarker
+where
+    In: Copy,
+    A: System<In = In, Out = bool>,
+    B: System<In = In, Out = bool>,
+{
+    type In = In;
+    type Out = bool;
+
+    fn combine(
+        input: Self::In,
+        world: &World,
+        a: impl FnOnce(<A as System>::In, &World) -> <A as System>::Out,
+        b: impl FnOnce(<B as System>::In, &World) -> <B as System>::Out,
+    ) -> Self::Out {
+        a(input, world) && b(input, world)
+    }
+
+    fn combine_exclusive(
+        input: Self::In,
+        world: &mut World,
+        a: impl FnOnce(<A as System>::In, &mut World) -> <A as System>::Out,
+        b: impl FnOnce(<B as System>::In, &mut World) -> <B as System>::Out,
+    ) -> Self::Out {
+        a(input, world) && b(input, world)
+    }
+}
+
+#[doc(hidden)]
+pub struct OrElseMarker;
+
+impl<In, A, B> Combine<A, B> for OrElseMarker
+where
+    In: Copy,
+    A: System<In = In, Out = bool>,
+    B: System<In = In, Out = bool>,
+{
+    type In = In;
+    type Out = bool;
+
+    fn combine(
+        input: Self::In,
+        world: &World,
+        a: impl FnOnce(<A as System>::In, &World) -> <A as System>::Out,
+        b: impl FnOnce(<B as System>::In, &World) -> <B as System>::Out,
+    ) -> Self::Out {
+        a(input, world) || b(input, world)
+    }
+
+    fn combine_exclusive(
+        input: Self::In,
+        world: &mut World,
+        a: impl FnOnce(<A as System>::In, &mut World) -> <A as System>::Out,
+        b: impl FnOnce(<B as System>::In, &mut World) -> <B as System>::Out,
+    ) -> Self::Out {
+        a(input, world) || b(input, world)
+    }
+}

crates/bevy_ecs/src/system/combinator.rs

@@ -0,0 +1,228 @@
+use std::{borrow::Cow, marker::PhantomData};
+
+use crate::{
+    archetype::ArchetypeComponentId, component::ComponentId, prelude::World, query::Access,
+};
+
+use super::{ReadOnlySystem, System};
+
+/// Customizes the behavior of a [`CombinatorSystem`].
+///
+/// # Examples
+///
+/// ```
+/// use bevy_ecs::prelude::*;
+/// use bevy_ecs::system::{CombinatorSystem, Combine};
+///
+/// // A system combinator that performs an exclusive-or (XOR)
+/// // operation on the output of two systems.
+/// pub type Xor<A, B> = CombinatorSystem<A, B, XorMarker>;
+///
+/// // This struct is used to customize the behavior of our combinator.
+/// pub struct XorMarker;
+///
+/// impl<A, B> Combine for XorMarker
+///     where A: System<In = (), Out = bool>,
+///     where B: System<In = (), Out = bool>,
+/// {
+///     type In = ();
+///     type Out = bool;
+///
+///     fn combine(
+///         _input: Self::In,
+///         world: &World,
+///         a: impl FnOnce(A::In, &World) -> A::Out,
+///         b: impl FnOnce(B::In, &World) -> B::Out,
+///     ) -> Self::Out {
+///         a((), world) ^ b((), world)
+///     }
+///
+///     fn combine_exclusive(
+///         _input: Self::In,
+///         world: &mut World,
+///         a: impl FnOnce(A::In, &mut World) -> A::Out,
+///         b: impl FnOnce(B::In, &mut World) -> B::Out,
+///     ) -> Self::Out {
+///         a((), world) ^ b((), world)
+///     }
+/// }
+///
+/// # #[derive(Resource) struct A(u32);
+/// # #[derive(Resource) struct B(u32);
+/// # #[derive(Resource, Default) struct RanFlag(bool);
+/// # let mut world = World::new();
+/// # world.init_resource::<RanFlag>();
+/// #
+/// # let mut app = Schedule::new();
+/// app.add_system(my_system.run_if(Xor::new(
+///     state_equals(A(1)),
+///     state_equals(B(1)),
+///     // The name of the combined system.
+///     Cow::Borrowed("a ^ b"),
+/// )));
+/// # fn my_system(mut flag: ResMut<RanFlag>) { flag.0 = true; }
+/// #
+/// # world.insert_resource(A(0));
+/// # world.insert_resoruce(B(0));
+/// # schedule.run(&mut world);
+/// # // Neither condition passes, so the system does not run.
+/// # assert!(!world.resource::<RanFlag>().0);
+/// #
+/// # world.insert_resource(A(1));
+/// # schedule.run(&mut world);
+/// # // Only the first condition passes, so the system runs.
+/// # assert!(world.resource::<RanFlag>().0);
+/// # world.resource_mut::<RanFlag>().0 = false;
+/// #
+/// # world.insert_resource(B(1));
+/// # schedule.run(&mut world);
+/// # // Both conditions pass, so the system does not run.
+/// # assert!(!world.resource::<RanFlag>().0);
+/// #
+/// # world.insert_resource(A(0));
+/// # schedule.run(&mut world);
+/// # // Only the second condition passes, so the system runs.
+/// # assert!(world.resource::<RanFlag>().0);
+/// # world.resource_mut::<RanFlag>().0 = false;
+/// ```
+pub trait Combine<A: System, B: System> {
+    type In;
+    type Out;
+
+    fn combine(
+        input: Self::In,
+        world: &World,
+        a: impl FnOnce(A::In, &World) -> A::Out,
+        b: impl FnOnce(B::In, &World) -> B::Out,
+    ) -> Self::Out;
+
+    fn combine_exclusive(
+        input: Self::In,
+        world: &mut World,
+        a: impl FnOnce(A::In, &mut World) -> A::Out,
+        b: impl FnOnce(B::In, &mut World) -> B::Out,
+    ) -> Self::Out;
+}
+
+/// A [`System`] defined by combining two other systems.
+/// The behavior of this combinator is specified by implementing the [`Combine`] trait.
+pub struct CombinatorSystem<Func, A, B> {
+    _marker: PhantomData<fn() -> Func>,
+    a: A,
+    b: B,
+    name: Cow<'static, str>,
+    component_access: Access<ComponentId>,
+    archetype_component_access: Access<ArchetypeComponentId>,
+}
+
+impl<Func, A, B> CombinatorSystem<Func, A, B> {
+    pub const fn new(a: A, b: B, name: Cow<'static, str>) -> Self {
+        Self {
+            _marker: PhantomData,
+            a,
+            b,
+            name,
+            component_access: Access::new(),
+            archetype_component_access: Access::new(),
+        }
+    }
+}
+
+impl<A, B, Func> System for CombinatorSystem<Func, A, B>
+where
+    Func: Combine<A, B> + 'static,
+    A: System,
+    B: System,
+{
+    type In = Func::In;
+    type Out = Func::Out;
+
+    fn name(&self) -> Cow<'static, str> {
+        self.name.clone()
+    }
+
+    fn type_id(&self) -> std::any::TypeId {
+        std::any::TypeId::of::<Self>()
+    }
+
+    fn component_access(&self) -> &crate::query::Access<crate::component::ComponentId> {
+        &self.component_access
+    }
+
+    fn archetype_component_access(
+        &self,
+    ) -> &crate::query::Access<crate::archetype::ArchetypeComponentId> {
+        &self.archetype_component_access
+    }
+
+    fn is_send(&self) -> bool {
+        self.a.is_send() && self.b.is_send()
+    }
+
+    fn is_exclusive(&self) -> bool {
+        self.a.is_exclusive() || self.b.is_exclusive()
+    }
+
+    unsafe fn run_unsafe(&mut self, input: Self::In, world: &crate::prelude::World) -> Self::Out {
+        Func::combine(
+            input,
+            world,
+            |input, w| self.a.run_unsafe(input, w),
+            |input, w| self.b.run_unsafe(input, w),
+        )
+    }
+
+    fn run(&mut self, input: Self::In, world: &mut World) -> Self::Out {
+        Func::combine_exclusive(
+            input,
+            world,
+            |input, w| self.a.run(input, w),
+            |input, w| self.b.run(input, w),
+        )
+    }
+
+    fn apply_buffers(&mut self, world: &mut crate::prelude::World) {
+        self.a.apply_buffers(world);
+        self.b.apply_buffers(world);
+    }
+
+    fn initialize(&mut self, world: &mut crate::prelude::World) {
+        self.a.initialize(world);
+        self.b.initialize(world);
+        self.component_access.extend(self.a.component_access());
+        self.component_access.extend(self.b.component_access());
+    }
+
+    fn update_archetype_component_access(&mut self, world: &crate::prelude::World) {
+        self.a.update_archetype_component_access(world);
+        self.b.update_archetype_component_access(world);
+
+        self.archetype_component_access
+            .extend(self.a.archetype_component_access());
+        self.archetype_component_access
+            .extend(self.b.archetype_component_access());
+    }
+
+    fn check_change_tick(&mut self, change_tick: u32) {
+        self.a.check_change_tick(change_tick);
+        self.b.check_change_tick(change_tick);
+    }
+
+    fn get_last_change_tick(&self) -> u32 {
+        self.a.get_last_change_tick()
+    }
+
+    fn set_last_change_tick(&mut self, last_change_tick: u32) {
+        self.a.set_last_change_tick(last_change_tick);
+        self.b.set_last_change_tick(last_change_tick);
+    }
+}
+
+/// SAFETY: Both systems are read-only, so any system created by combining them will only read from the world.
+unsafe impl<A, B, Func> ReadOnlySystem for CombinatorSystem<Func, A, B>
+where
+    Func: Combine<A, B> + 'static,
+    A: ReadOnlySystem,
+    B: ReadOnlySystem,
+{
+}

crates/bevy_ecs/src/system/mod.rs

@@ -98,6 +98,7 @@
 //! - All tuples between 1 to 16 elements where each element implements [`SystemParam`]
 //! - [`()` (unit primitive type)](https://doc.rust-lang.org/stable/std/primitive.unit.html)
 
+mod combinator;
 mod commands;
 mod exclusive_function_system;
 mod exclusive_system_param;
@@ -108,6 +109,7 @@ mod system;
 mod system_param;
 mod system_piping;
 
+pub use combinator::*;
 pub use commands::*;
 pub use exclusive_function_system::*;
 pub use exclusive_system_param::*;