Simplify system building with BuilderSystem

Author: ecoskey

crates/bevy_ecs/src/system/builder.rs

@@ -1,23 +1,27 @@
 use alloc::{boxed::Box, vec::Vec};
 use bevy_platform::cell::SyncCell;
+use bevy_utils::prelude::DebugName;
 use variadics_please::all_tuples;
 
 use crate::{
+    change_detection::{CheckChangeTicks, Tick},
     prelude::QueryBuilder,
-    query::{QueryData, QueryFilter, QueryState},
+    query::{FilteredAccessSet, QueryData, QueryFilter, QueryState},
     resource::Resource,
     system::{
-        DynSystemParam, DynSystemParamState, If, Local, ParamSet, Query, SystemParam,
-        SystemParamValidationError,
+        DynSystemParam, DynSystemParamState, FromInput, FunctionSystem, If, IntoResult, IntoSystem,
+        Local, ParamSet, Query, ReadOnlySystem, System, SystemInput, SystemMeta, SystemParam,
+        SystemParamFunction, SystemParamValidationError,
     },
     world::{
-        FilteredResources, FilteredResourcesBuilder, FilteredResourcesMut,
-        FilteredResourcesMutBuilder, FromWorld, World,
+        unsafe_world_cell::UnsafeWorldCell, DeferredWorld, FilteredResources,
+        FilteredResourcesBuilder, FilteredResourcesMut, FilteredResourcesMutBuilder, FromWorld,
+        World,
     },
 };
-use core::fmt::Debug;
+use core::{fmt::Debug, marker::PhantomData, mem};
 
-use super::{Res, ResMut, SystemState};
+use super::{Res, ResMut, RunSystemError, SystemState, SystemStateFlags};
 
 /// A builder that can create a [`SystemParam`].
 ///
@@ -119,6 +123,17 @@ pub unsafe trait SystemParamBuilder<P: SystemParam>: Sized {
     fn build_state(self, world: &mut World) -> SystemState<P> {
         SystemState::from_builder(world, self)
     }
+
+    /// Create a [`System`] from a [`SystemParamBuilder`]
+    fn build_system<Marker, In, Out, Func>(
+        self,
+        func: Func,
+    ) -> IntoBuilderSystem<Marker, In, Out, Func, Self>
+    where
+        Func: SystemParamFunction<Marker, Param = P>,
+    {
+        IntoBuilderSystem::new(self, func)
+    }
 }
 
 /// A [`SystemParamBuilder`] for any [`SystemParam`] that uses its default initialization.
@@ -204,6 +219,254 @@ impl ParamBuilder {
     }
 }
 
+/// A marker type used to distinguish builder systems from plain function systems.
+#[doc(hidden)]
+pub struct IsBuilderSystem;
+
+/// An [`IntoSystem`] creating an instance of [`BuilderSystem`]
+pub struct IntoBuilderSystem<Marker, In, Out, Func, Builder>
+where
+    Func: SystemParamFunction<Marker>,
+    Builder: SystemParamBuilder<Func::Param>,
+{
+    builder: Builder,
+    func: Func,
+    _data: PhantomData<fn(In) -> (Marker, Out)>,
+}
+
+impl<Marker, In, Out, Func, Builder> IntoBuilderSystem<Marker, In, Out, Func, Builder>
+where
+    Func: SystemParamFunction<Marker>,
+    Builder: SystemParamBuilder<Func::Param>,
+{
+    /// Returns a new [`IntoBuilderSystem`] given a system param builder and system function
+    pub fn new(builder: Builder, func: Func) -> Self {
+        Self {
+            builder,
+            func,
+            _data: PhantomData,
+        }
+    }
+}
+
+impl<Marker, In, Out, Func, Builder> IntoSystem<In, Out, (IsBuilderSystem, Marker)>
+    for IntoBuilderSystem<Marker, In, Out, Func, Builder>
+where
+    Marker: 'static,
+    In: SystemInput + 'static,
+    Out: 'static,
+    Func: SystemParamFunction<Marker, In: FromInput<In>, Out: IntoResult<Out>>,
+    Builder: SystemParamBuilder<Func::Param> + Send + Sync + 'static,
+{
+    type System = BuilderSystem<Marker, In, Out, Func, Builder>;
+
+    fn into_system(this: Self) -> Self::System {
+        BuilderSystem::new(this.builder, this.func)
+    }
+}
+
+/// A [`System`] created from a [`SystemParamBuilder`] whose state is not
+/// initialized until the first run.
+pub struct BuilderSystem<Marker, In, Out, Func, Builder>
+where
+    Func: SystemParamFunction<Marker>,
+    Builder: SystemParamBuilder<Func::Param>,
+{
+    inner: BuilderSystemInner<Marker, In, Out, Func, Builder>,
+}
+
+impl<Marker, In, Out, Func, Builder> BuilderSystem<Marker, In, Out, Func, Builder>
+where
+    Func: SystemParamFunction<Marker>,
+    Builder: SystemParamBuilder<Func::Param>,
+{
+    /// Returns a new `BuilderSystem` given a system param builder and a system function
+    pub fn new(builder: Builder, func: Func) -> Self {
+        Self {
+            inner: BuilderSystemInner::Uninitialized {
+                builder,
+                func,
+                meta: SystemMeta::new::<Func>(),
+            },
+        }
+    }
+}
+
+enum BuilderSystemInner<Marker, In, Out, Func, Builder>
+where
+    Func: SystemParamFunction<Marker>,
+    Builder: SystemParamBuilder<Func::Param>,
+{
+    /// A properly initialized system whose state has been constructed
+    Initialized {
+        system: FunctionSystem<Marker, In, Out, Func>,
+    },
+    /// An unititialized system, whose state hasn't been constructed from
+    /// the param builder yet
+    Uninitialized {
+        builder: Builder,
+        func: Func,
+        meta: SystemMeta,
+    },
+    // This only exists as a variant to use with `mem::replace` in `initialize`.
+    // If this state is ever observed outside `initialize`, then a `panic!`
+    // interrupted initialization, leaving this system in an invalid state.
+    Invalid,
+}
+
+impl<Marker, In, Out, Func, Builder> System for BuilderSystem<Marker, In, Out, Func, Builder>
+where
+    Marker: 'static,
+    In: SystemInput + 'static,
+    Out: 'static,
+    Func: SystemParamFunction<Marker, In: FromInput<In>, Out: IntoResult<Out>>,
+    Builder: SystemParamBuilder<Func::Param> + Send + Sync + 'static,
+{
+    type In = In;
+
+    type Out = Out;
+
+    #[inline]
+    fn name(&self) -> DebugName {
+        match &self.inner {
+            BuilderSystemInner::Initialized { system } => system.name(),
+            BuilderSystemInner::Uninitialized { meta, .. } => meta.name().clone(),
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[inline]
+    fn flags(&self) -> SystemStateFlags {
+        match &self.inner {
+            BuilderSystemInner::Initialized { system, .. } => system.flags(),
+            BuilderSystemInner::Uninitialized { meta, .. } => meta.flags(),
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[inline]
+    unsafe fn run_unsafe(
+        &mut self,
+        input: super::SystemIn<'_, Self>,
+        world: UnsafeWorldCell,
+    ) -> Result<Self::Out, RunSystemError> {
+        match &mut self.inner {
+            // SAFETY: requirements upheld by the caller.
+            BuilderSystemInner::Initialized { system, .. } => unsafe {
+                system.run_unsafe(input, world)
+            },
+            BuilderSystemInner::Uninitialized { .. } => panic!(
+                "BuilderSystem {} was not initialized before calling run_unsafe.",
+                self.name()
+            ),
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[cfg(feature = "hotpatching")]
+    #[inline]
+    fn refresh_hotpatch(&mut self) {
+        match &mut self.inner {
+            BuilderSystemInner::Initialized { system, .. } => system.refresh_hotpatch(),
+            BuilderSystemInner::Uninitialized { .. } => {}
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[inline]
+    fn apply_deferred(&mut self, world: &mut World) {
+        match &mut self.inner {
+            BuilderSystemInner::Initialized { system, .. } => system.apply_deferred(world),
+            BuilderSystemInner::Uninitialized { .. } => {}
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[inline]
+    fn queue_deferred(&mut self, world: DeferredWorld) {
+        match &mut self.inner {
+            BuilderSystemInner::Initialized { system, .. } => system.queue_deferred(world),
+            BuilderSystemInner::Uninitialized { .. } => {}
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[inline]
+    unsafe fn validate_param_unsafe(
+        &mut self,
+        world: UnsafeWorldCell,
+    ) -> Result<(), SystemParamValidationError> {
+        match &mut self.inner {
+            // SAFETY: requirements upheld by the caller.
+            BuilderSystemInner::Initialized { system, .. } => unsafe {
+                system.validate_param_unsafe(world)
+            },
+            BuilderSystemInner::Uninitialized { .. } => Ok(()),
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[inline]
+    fn initialize(&mut self, world: &mut World) -> FilteredAccessSet {
+        let inner = mem::replace(&mut self.inner, BuilderSystemInner::Invalid);
+        match inner {
+            BuilderSystemInner::Initialized { mut system } => {
+                let access = system.initialize(world);
+                self.inner = BuilderSystemInner::Initialized { system };
+                access
+            }
+            BuilderSystemInner::Uninitialized { builder, func, .. } => {
+                let mut system = builder.build_state(world).build_any_system(func);
+                let access = system.initialize(world);
+                self.inner = BuilderSystemInner::Initialized { system };
+                access
+            }
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[inline]
+    fn check_change_tick(&mut self, check: CheckChangeTicks) {
+        match &mut self.inner {
+            BuilderSystemInner::Initialized { system, .. } => system.check_change_tick(check),
+            BuilderSystemInner::Uninitialized { .. } => {}
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[inline]
+    fn get_last_run(&self) -> Tick {
+        match &self.inner {
+            BuilderSystemInner::Initialized { system, .. } => system.get_last_run(),
+            BuilderSystemInner::Uninitialized { meta, .. } => meta.get_last_run(),
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+
+    #[inline]
+    fn set_last_run(&mut self, last_run: Tick) {
+        match &mut self.inner {
+            BuilderSystemInner::Initialized { system, .. } => system.set_last_run(last_run),
+            BuilderSystemInner::Uninitialized { meta, .. } => meta.set_last_run(last_run),
+            BuilderSystemInner::Invalid => unreachable!(),
+        }
+    }
+}
+
+// SAFETY: if the wrapped system is read-only, so is this one
+unsafe impl<Marker, In, Out, Func, Builder> ReadOnlySystem
+    for BuilderSystem<Marker, In, Out, Func, Builder>
+where
+    Marker: 'static,
+    In: SystemInput + 'static,
+    Out: 'static,
+    Func: SystemParamFunction<Marker, In: FromInput<In>, Out: IntoResult<Out>>,
+    Builder: SystemParamBuilder<Func::Param> + Send + Sync + 'static,
+    // the important bound
+    FunctionSystem<Marker, In, Out, Func>: ReadOnlySystem,
+{
+}
+
 // SAFETY: Any `QueryState<D, F>` for the correct world is valid for `Query::State`,
 // and we check the world during `build`.
 unsafe impl<'w, 's, D: QueryData + 'static, F: QueryFilter + 'static>

crates/bevy_ecs/src/system/function_system.rs

@@ -64,6 +64,11 @@ impl SystemMeta {
         &self.name
     }
 
+    /// Returns the system's state flags
+    pub fn flags(&self) -> SystemStateFlags {
+        self.flags
+    }
+
     /// Sets the name of this system.
     ///
     /// Useful to give closure systems more readable and unique names for debugging and tracing.
@@ -79,6 +84,18 @@ impl SystemMeta {
         self.name = new_name.into();
     }
 
+    /// Gets the last time this system was run.
+    #[inline]
+    pub fn get_last_run(&self) -> Tick {
+        self.last_run
+    }
+
+    /// Sets the last time this system was run.
+    #[inline]
+    pub fn set_last_run(&mut self, last_run: Tick) {
+        self.last_run = last_run;
+    }
+
     /// Returns true if the system is [`Send`].
     #[inline]
     pub fn is_send(&self) -> bool {