๐Ÿ“ฆ Zeenobit / moonshine_util

๐Ÿ“„ defer.rs ยท 142 lines
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142use std::marker::PhantomData;

use bevy_app::prelude::*;
use bevy_ecs::prelude::*;
use bevy_ecs::schedule::ScheduleLabel;
use bevy_ecs::system::SystemId;
use bevy_log::prelude::*;

use crate::Static;

pub struct DefaultDeferredSystemsPlugin;

impl Plugin for DefaultDeferredSystemsPlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(First, run_deferred_systems::<First>)
            .add_systems(PreUpdate, run_deferred_systems::<PreUpdate>)
            .add_systems(Update, run_deferred_systems::<Update>)
            .add_systems(PostUpdate, run_deferred_systems::<PostUpdate>)
            .add_systems(Last, run_deferred_systems::<Last>);
    }
}

pub trait RunDeferredSystem {
    fn run_deferred_system<S: ScheduleLabel, M>(
        &mut self,
        schedule: S,
        system: impl 'static + IntoSystem<(), (), M>,
    );

    fn run_deferred_system_with<S: ScheduleLabel, I: Static, M>(
        &mut self,
        schedule: S,
        system: impl 'static + IntoSystem<In<I>, (), M>,
        input: I,
    );
}

impl RunDeferredSystem for World {
    fn run_deferred_system<S: ScheduleLabel, M>(
        &mut self,
        _schedule: S,
        system: impl 'static + IntoSystem<(), (), M>,
    ) {
        let system = self.register_system_cached(system);
        self.get_resource_or_init::<DeferredSystems<S>>()
            .0
            .push(Box::new(DeferredSystem(system)));
    }

    fn run_deferred_system_with<S: ScheduleLabel, I: Static, M>(
        &mut self,
        _schedule: S,
        system: impl 'static + IntoSystem<In<I>, (), M>,
        input: I,
    ) {
        let system = self.register_system_cached(system);
        self.get_resource_or_init::<DeferredSystems<S>>()
            .0
            .push(Box::new(DeferredSystemWith(system, input)));
    }
}

#[derive(Resource)]
struct DeferredSystems<S: ScheduleLabel>(Vec<Box<dyn AnyDeferredSystem>>, PhantomData<S>);

impl<S: ScheduleLabel> Default for DeferredSystems<S> {
    fn default() -> Self {
        Self(Vec::new(), PhantomData)
    }
}

impl<S: ScheduleLabel> DeferredSystems<S> {
    fn take(&mut self) -> Self {
        Self(self.0.drain(..).collect(), PhantomData)
    }

    fn run(self, world: &mut World) {
        for system in self.0 {
            system.run(world);
        }
    }
}

trait AnyDeferredSystem: Static {
    fn run(self: Box<Self>, world: &mut World);
}

struct DeferredSystem(SystemId);

impl AnyDeferredSystem for DeferredSystem {
    fn run(self: Box<Self>, world: &mut World) {
        if let Err(why) = world.run_system(self.0) {
            error!("deferred system error: {why}");
        }
    }
}

struct DeferredSystemWith<I: Static>(SystemId<In<I>>, I);

impl<I: Static> AnyDeferredSystem for DeferredSystemWith<I> {
    fn run(self: Box<Self>, world: &mut World) {
        if let Err(why) = world.run_system_with(self.0, self.1) {
            error!("deferred system error: {why}");
        }
    }
}

pub fn run_deferred_systems<S: ScheduleLabel>(world: &mut World) {
    let Some(mut systems) = world.get_resource_mut::<DeferredSystems<S>>() else {
        return;
    };

    systems.take().run(world);
}

#[test]
fn test_deferred_system() {
    use bevy::prelude::*;
    use bevy::MinimalPlugins;

    #[derive(Resource)]
    struct Success;

    let mut app = App::new();
    app.add_plugins((MinimalPlugins, DefaultDeferredSystemsPlugin));

    app.world_mut()
        .run_deferred_system(Update, |mut commands: Commands| {
            commands.insert_resource(Success)
        });

    app.world_mut().flush(); // Should be redundant, but just to be sure ...
    assert!(!app.world().contains_resource::<Success>());

    app.update();
    assert!(app.world_mut().remove_resource::<Success>().is_some());

    // Ensure the system does not run again
    app.update();
    assert!(!app.world().contains_resource::<Success>());
}