Skip to content

Commit

Permalink
Cache systems by S instead of S::System (#16694)
Browse files Browse the repository at this point in the history
# Objective

- Fixes the issue described in this comment:
#16680 (comment).

## Solution

- Cache one-shot systems by `S: IntoSystem` (which is const-asserted to
be a ZST) rather than `S::System`.

## Testing

Added a new unit test named `cached_system_into_same_system_type` to
`system_registry.rs`.

---

## Migration Guide

The `CachedSystemId` resource has been changed:

```rust
// Before:
let cached_id = CachedSystemId::<S::System>(id);
assert!(id == cached_id.0);

// After:
let cached_id = CachedSystemId::<S>::new(id);
assert!(id == SystemId::from_entity(cached_id.entity));
```
  • Loading branch information
benfrankel authored Mar 4, 2025
1 parent f3db44b commit 0a841ba
Showing 1 changed file with 62 additions and 11 deletions.
73 changes: 62 additions & 11 deletions crates/bevy_ecs/src/system/system_registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use crate::reflect::ReflectComponent;
use crate::{
change_detection::Mut,
entity::Entity,
system::{input::SystemInput, BoxedSystem, IntoSystem, System},
system::{input::SystemInput, BoxedSystem, IntoSystem},
world::World,
};
use alloc::boxed::Box;
Expand Down Expand Up @@ -126,7 +126,21 @@ impl<I: SystemInput, O> core::fmt::Debug for SystemId<I, O> {
///
/// This resource is inserted by [`World::register_system_cached`].
#[derive(Resource)]
pub struct CachedSystemId<S: System>(pub SystemId<S::In, S::Out>);
pub struct CachedSystemId<S> {
/// The cached `SystemId` as an `Entity`.
pub entity: Entity,
_marker: PhantomData<fn() -> S>,
}

impl<S> CachedSystemId<S> {
/// Creates a new `CachedSystemId` struct given a `SystemId`.
pub fn new<I: SystemInput, O>(id: SystemId<I, O>) -> Self {
Self {
entity: id.entity(),
_marker: PhantomData,
}
}
}

impl World {
/// Registers a system and returns a [`SystemId`] so it can later be called by [`World::run_system`].
Expand Down Expand Up @@ -154,7 +168,7 @@ impl World {
/// Similar to [`Self::register_system`], but allows passing in a [`BoxedSystem`].
///
/// This is useful if the [`IntoSystem`] implementor has already been turned into a
/// [`System`] trait object and put in a [`Box`].
/// [`System`](crate::system::System) trait object and put in a [`Box`].
pub fn register_boxed_system<I, O>(&mut self, system: BoxedSystem<I, O>) -> SystemId<I, O>
where
I: SystemInput + 'static,
Expand Down Expand Up @@ -389,23 +403,23 @@ impl World {
);
}

if !self.contains_resource::<CachedSystemId<S::System>>() {
if !self.contains_resource::<CachedSystemId<S>>() {
let id = self.register_system(system);
self.insert_resource(CachedSystemId::<S::System>(id));
self.insert_resource(CachedSystemId::<S>::new(id));
return id;
}

self.resource_scope(|world, mut id: Mut<CachedSystemId<S::System>>| {
if let Ok(mut entity) = world.get_entity_mut(id.0.entity()) {
self.resource_scope(|world, mut id: Mut<CachedSystemId<S>>| {
if let Ok(mut entity) = world.get_entity_mut(id.entity) {
if !entity.contains::<RegisteredSystem<I, O>>() {
entity.insert(RegisteredSystem::new(Box::new(IntoSystem::into_system(
system,
))));
}
} else {
id.0 = world.register_system(system);
id.entity = world.register_system(system).entity();
}
id.0
SystemId::from_entity(id.entity)
})
}

Expand All @@ -422,9 +436,9 @@ impl World {
S: IntoSystem<I, O, M> + 'static,
{
let id = self
.remove_resource::<CachedSystemId<S::System>>()
.remove_resource::<CachedSystemId<S>>()
.ok_or(RegisteredSystemError::SystemNotCached)?;
self.unregister_system(id.0)
self.unregister_system(SystemId::<I, O>::from_entity(id.entity))
}

/// Runs a cached system, registering it if necessary.
Expand Down Expand Up @@ -754,6 +768,43 @@ mod tests {
assert!(matches!(output, Ok(8)));
}

#[test]
fn cached_system_into_same_system_type() {
use crate::result::Result;

struct Foo;
impl IntoSystem<(), Result<()>, ()> for Foo {
type System = ApplyDeferred;
fn into_system(_: Self) -> Self::System {
ApplyDeferred
}
}

struct Bar;
impl IntoSystem<(), Result<()>, ()> for Bar {
type System = ApplyDeferred;
fn into_system(_: Self) -> Self::System {
ApplyDeferred
}
}

let mut world = World::new();
let foo1 = world.register_system_cached(Foo);
let foo2 = world.register_system_cached(Foo);
let bar1 = world.register_system_cached(Bar);
let bar2 = world.register_system_cached(Bar);

// The `S: IntoSystem` types are different, so they should be cached
// as separate systems, even though the `<S as IntoSystem>::System`
// types / values are the same (`ApplyDeferred`).
assert_ne!(foo1, bar1);

// But if the `S: IntoSystem` types are the same, they'll be cached
// as the same system.
assert_eq!(foo1, foo2);
assert_eq!(bar1, bar2);
}

#[test]
fn system_with_input_ref() {
fn with_ref(InRef(input): InRef<u8>, mut counter: ResMut<Counter>) {
Expand Down

0 comments on commit 0a841ba

Please sign in to comment.