Skip to main content

moonshine_save/
save.rs

1use std::any::TypeId;
2use std::io::{self, Write};
3use std::marker::PhantomData;
4use std::path::PathBuf;
5
6use bevy_ecs::entity::EntityHashSet;
7use bevy_ecs::prelude::*;
8use bevy_ecs::query::QueryFilter;
9use bevy_log::prelude::*;
10use bevy_world_serialization::{DynamicWorld, DynamicWorldBuilder, WorldFilter};
11
12use moonshine_util::event::{OnSingle, SingleEvent, TriggerSingle};
13use moonshine_util::Static;
14use thiserror::Error;
15
16use crate::{MapComponent, SceneMapper};
17
18/// A [`Component`] which marks its [`Entity`] to be saved.
19#[derive(Component, Default, Debug, Clone)]
20pub struct Save;
21
22/// A trait used to trigger a [`SaveEvent`] via [`Commands`] or [`World`].
23pub trait TriggerSave {
24    /// Triggers the given [`SaveEvent`].
25    #[doc(alias = "trigger_single")]
26    fn trigger_save(self, event: impl SaveEvent);
27}
28
29impl TriggerSave for &mut Commands<'_, '_> {
30    fn trigger_save(self, event: impl SaveEvent) {
31        self.trigger_single(event);
32    }
33}
34
35impl TriggerSave for &mut World {
36    fn trigger_save(self, event: impl SaveEvent) {
37        self.trigger_single(event);
38    }
39}
40
41/// A [`SingleEvent`] which starts the save process with the given parameters.
42///
43/// See also:
44/// - [`trigger_save`](TriggerSave::trigger_save)
45/// - [`trigger_single`](TriggerSingle::trigger_single)
46/// - [`SaveWorld`]
47pub trait SaveEvent: SingleEvent {
48    /// A [`QueryFilter`] used as the initial filter for selecting saved entities.
49    type SaveFilter: QueryFilter;
50
51    /// Return `true` if the given [`Entity`] should be saved.
52    fn filter_entity(&self, _entity: EntityRef) -> bool {
53        true
54    }
55
56    /// Called once before the save process starts.
57    ///
58    /// This is useful if you want to modify the world just before saving.
59    fn before_save(&mut self, _world: &mut World) {}
60
61    /// Called once before serialization.
62    ///
63    /// This is useful to undo any modifications done before saving.
64    fn before_serialize(&mut self, _world: &mut World, _entities: &[Entity]) {}
65
66    /// Returns a [`WorldFilter`] for selecting which components should be saved.
67    fn component_filter(&mut self) -> WorldFilter {
68        WorldFilter::allow_all()
69    }
70
71    /// Returns a [`WorldFilter`] for selecting which resources should be saved.
72    fn resource_filter(&mut self) -> WorldFilter {
73        WorldFilter::deny_all()
74    }
75
76    /// Called once after serialization.
77    ///
78    /// This is useful if you would like to do any post-processing of the [`Saved`] data *before* [`OnSave`] is triggered.
79    fn after_save(&mut self, _world: &mut World, _result: &SaveResult) {}
80
81    /// Returns the [`SaveOutput`] of the save process.
82    fn output(&mut self) -> SaveOutput;
83}
84
85/// A generic [`SaveEvent`] which can be used to save the [`World`].
86pub struct SaveWorld<F: QueryFilter = DefaultSaveFilter> {
87    /// A filter for selecting which entities should be saved.
88    ///
89    /// By default, all entities are selected.
90    pub entities: EntityFilter,
91    /// A filter for selecting which resources should be saved.
92    ///
93    /// By default, no resources are selected. Most Bevy resources are not safely serializable.
94    pub resources: WorldFilter,
95    /// A filter for selecting which components should be saved.
96    ///
97    /// By default, all serializable components are selected.
98    pub components: WorldFilter,
99    /// A mapper for transforming components during the save process.
100    ///
101    /// See [`MapComponent`] for more information.
102    pub mapper: SceneMapper,
103    /// Output of the saved world.
104    pub output: SaveOutput,
105    #[doc(hidden)]
106    pub filter: PhantomData<F>,
107}
108
109impl<F: QueryFilter> SaveWorld<F> {
110    /// Creates a new [`SaveWorld`] event with the given [`SaveOutput`].
111    pub fn new(output: SaveOutput) -> Self {
112        Self {
113            entities: EntityFilter::allow_all(),
114            resources: WorldFilter::deny_all(),
115            components: WorldFilter::allow_all(),
116            mapper: SceneMapper::default(),
117            output,
118            filter: PhantomData,
119        }
120    }
121
122    /// Creates a new [`SaveWorld`] event which saves entities matching the
123    /// given [`QueryFilter`] into a file at the given path.
124    pub fn into_file(path: impl Into<PathBuf>) -> Self {
125        Self {
126            entities: EntityFilter::allow_all(),
127            resources: WorldFilter::deny_all(),
128            components: WorldFilter::allow_all(),
129            mapper: SceneMapper::default(),
130            output: SaveOutput::file(path),
131            filter: PhantomData,
132        }
133    }
134
135    /// Creates a new [`SaveWorld`] event which saves entities matching the
136    /// given [`QueryFilter`] into a [`Write`] stream.
137    pub fn into_stream(stream: impl SaveStream) -> Self {
138        Self {
139            entities: EntityFilter::allow_all(),
140            resources: WorldFilter::deny_all(),
141            components: WorldFilter::allow_all(),
142            mapper: SceneMapper::default(),
143            output: SaveOutput::stream(stream),
144            filter: PhantomData,
145        }
146    }
147
148    /// Includes the given [`Resource`] in the save data.
149    pub fn include_resource<R: Resource>(mut self) -> Self {
150        self.resources = self.resources.allow::<R>();
151        self
152    }
153
154    /// Includes the given [`Resource`] by its [`TypeId`] in the save data.
155    pub fn include_resource_by_id(mut self, type_id: TypeId) -> Self {
156        self.resources = self.resources.allow_by_id(type_id);
157        self
158    }
159
160    /// Excludes the given [`Component`] from the save data.
161    pub fn exclude_component<T: Component>(mut self) -> Self {
162        self.components = self.components.deny::<T>();
163        self
164    }
165
166    /// Excludes the given [`Component`] by its [`TypeId`] from the save data.
167    pub fn exclude_component_by_id(mut self, type_id: TypeId) -> Self {
168        self.components = self.components.deny_by_id(type_id);
169        self
170    }
171
172    /// Maps the given [`Component`] into another using a [component mapper](MapComponent) before saving.
173    pub fn map_component<T: Component>(mut self, m: impl MapComponent<T>) -> Self {
174        self.mapper = self.mapper.map(m);
175        self
176    }
177}
178
179impl SaveWorld {
180    /// Creates a new [`SaveWorld`] event which saves default entities (with [`Save`])
181    /// into a file at the given path.
182    pub fn default_into_file(path: impl Into<PathBuf>) -> Self {
183        Self::into_file(path)
184    }
185
186    /// Creates a new [`SaveWorld`] event which saves default entities (with [`Save`])
187    /// into a [`Write`] stream.
188    pub fn default_into_stream(stream: impl SaveStream) -> Self {
189        Self::into_stream(stream)
190    }
191}
192
193impl SaveWorld<()> {
194    /// Creates a new [`SaveWorld`] event which saves all entities into a file at the given path.
195    pub fn all_into_file(path: impl Into<PathBuf>) -> Self {
196        Self::into_file(path)
197    }
198
199    /// Creates a new [`SaveWorld`] event which saves all entities into a [`Write`] stream.
200    pub fn all_into_stream(stream: impl SaveStream) -> Self {
201        Self::into_stream(stream)
202    }
203}
204
205impl<F: QueryFilter> SingleEvent for SaveWorld<F> where F: Static {}
206
207impl<F: QueryFilter> SaveEvent for SaveWorld<F>
208where
209    F: Static,
210{
211    type SaveFilter = F;
212
213    fn filter_entity(&self, entity: EntityRef) -> bool {
214        match &self.entities {
215            EntityFilter::Allow(allow) => allow.contains(&entity.id()),
216            EntityFilter::Block(block) => !block.contains(&entity.id()),
217        }
218    }
219
220    fn before_serialize(&mut self, world: &mut World, entities: &[Entity]) {
221        for entity in entities {
222            self.mapper.apply(world.entity_mut(*entity));
223        }
224    }
225
226    fn after_save(&mut self, world: &mut World, result: &SaveResult) {
227        let Ok(saved) = result else {
228            return;
229        };
230
231        for entity in saved.entities() {
232            self.mapper.undo(world.entity_mut(entity));
233        }
234    }
235
236    fn component_filter(&mut self) -> WorldFilter {
237        std::mem::replace(&mut self.components, WorldFilter::Unset)
238    }
239
240    fn resource_filter(&mut self) -> WorldFilter {
241        std::mem::replace(&mut self.resources, WorldFilter::Unset)
242    }
243
244    fn output(&mut self) -> SaveOutput {
245        self.output.consume().unwrap()
246    }
247}
248
249/// Filter used for the default [`SaveWorld`] event.
250/// This includes all entities with the [`Save`] component.
251pub type DefaultSaveFilter = With<Save>;
252
253/// Output of the save process.
254pub enum SaveOutput {
255    /// Save into a file at the given path.
256    File(PathBuf),
257    /// Save into a [`Write`] stream.
258    Stream(Box<dyn SaveStream>),
259    /// Drops the save data.
260    ///
261    /// This is useful if you would like to process the [`Saved`] data manually.
262    /// You can observe the [`OnSave`] event for post-processing logic.
263    Drop,
264    #[doc(hidden)]
265    Invalid,
266}
267
268impl SaveOutput {
269    /// Creates a new [`SaveOutput`] which saves into a file at the given path.
270    pub fn file(path: impl Into<PathBuf>) -> Self {
271        Self::File(path.into())
272    }
273
274    /// Creates a new [`SaveOutput`] which saves into a [`Write`] stream.
275    pub fn stream<S: SaveStream + 'static>(stream: S) -> Self {
276        Self::Stream(Box::new(stream))
277    }
278
279    /// Invalidates this [`SaveOutput`] and returns it if it was valid.
280    pub fn consume(&mut self) -> Option<SaveOutput> {
281        let output = std::mem::replace(self, SaveOutput::Invalid);
282        if let SaveOutput::Invalid = output {
283            return None;
284        }
285        Some(output)
286    }
287}
288
289/// A filter for selecting which [`Entity`]s within a [`World`].
290#[derive(Clone, Debug)]
291pub enum EntityFilter {
292    /// Select only the specified entities.
293    Allow(EntityHashSet),
294    /// Select all entities except the specified ones.
295    Block(EntityHashSet),
296}
297
298impl EntityFilter {
299    /// Creates a new [`EntityFilter`] which allows all entities.
300    pub fn allow_all() -> Self {
301        Self::Block(EntityHashSet::new())
302    }
303
304    /// Creates a new [`EntityFilter`] which allows only the specified entities.
305    pub fn allow(entities: impl IntoIterator<Item = Entity>) -> Self {
306        Self::Allow(entities.into_iter().collect())
307    }
308
309    /// Creates a new [`EntityFilter`] which blocks the specified entities.
310    pub fn block(entities: impl IntoIterator<Item = Entity>) -> Self {
311        Self::Block(entities.into_iter().collect())
312    }
313}
314
315impl Default for EntityFilter {
316    fn default() -> Self {
317        Self::allow_all()
318    }
319}
320
321/// Alias for a `'static` [`Write`] stream.
322pub trait SaveStream: Write
323where
324    Self: Static,
325{
326}
327
328impl<S: Write> SaveStream for S where S: Static {}
329
330/// An [`Event`] triggered at the end of the save process.
331///
332/// This event contains the saved [`World`] data as a [`DynamicWorld`].
333#[derive(Event)]
334pub struct Saved {
335    /// The saved [`DynamicWorld`] to be serialized.
336    pub world: DynamicWorld,
337}
338
339impl Saved {
340    /// Iterates over all the saved entities.
341    pub fn entities(&self) -> impl Iterator<Item = Entity> + '_ {
342        self.world.entities.iter().map(|de| de.entity)
343    }
344}
345
346#[doc(hidden)]
347#[deprecated(since = "0.5.2", note = "use `Saved` instead")]
348pub type OnSave = Saved;
349
350/// An error that may occur during the save process.
351#[derive(Error, Debug)]
352pub enum SaveError {
353    /// An error occurred while serializing the scene.
354    #[error("Failed to serialize world: {0}")]
355    Ron(ron::Error),
356    /// An error occurred while writing into [`SaveOutput`].
357    #[error("Failed to write world: {0}")]
358    Io(io::Error),
359}
360
361impl From<ron::Error> for SaveError {
362    fn from(e: ron::Error) -> Self {
363        Self::Ron(e)
364    }
365}
366
367impl From<io::Error> for SaveError {
368    fn from(e: io::Error) -> Self {
369        Self::Io(e)
370    }
371}
372
373/// [`Result`] of a [`SaveEvent`].
374pub type SaveResult = Result<Saved, SaveError>;
375
376/// An [`Observer`] which saved the world when a [`SaveWorld`] event is triggered.
377pub fn save_on_default_event(event: OnSingle<SaveWorld>, commands: Commands) {
378    save_on(event, commands);
379}
380
381/// An [`Observer`] which saved the world when the given [`SaveEvent`] is triggered.
382pub fn save_on<E: SaveEvent>(event: OnSingle<E>, mut commands: Commands) {
383    commands.queue_handled(SaveCommand(event.consume().unwrap()), |err, ctx| {
384        error!("save failed: {err:?} ({ctx})");
385    });
386}
387
388fn save_world<E: SaveEvent>(mut event: E, world: &mut World) -> SaveResult {
389    // Notify
390    event.before_save(world);
391
392    // Filter
393    let entities: Vec<_> = world
394        .query_filtered::<Entity, E::SaveFilter>()
395        .iter(world)
396        .filter(|entity| event.filter_entity(world.entity(*entity)))
397        .collect();
398
399    // Serialize
400    event.before_serialize(world, &entities);
401    let saved_world = {
402        let type_registry = world.resource::<AppTypeRegistry>().read();
403        DynamicWorldBuilder::from_world(world, &type_registry)
404            .with_component_filter(event.component_filter())
405            .with_resource_filter(event.resource_filter())
406            .extract_resources()
407            .extract_entities(entities.iter().copied())
408            .build()
409    };
410
411    // Write
412    let saved = match event.output() {
413        SaveOutput::File(path) => {
414            if let Some(parent) = path.parent() {
415                std::fs::create_dir_all(parent)?;
416            }
417
418            let type_registry = world.resource::<AppTypeRegistry>().read();
419            let data = saved_world.serialize(&type_registry)?;
420            std::fs::write(&path, data.as_bytes())?;
421            debug!("saved into file: {path:?}");
422            Saved { world: saved_world }
423        }
424        SaveOutput::Stream(mut stream) => {
425            let type_registry = world.resource::<AppTypeRegistry>().read();
426            let data = saved_world.serialize(&type_registry)?;
427            stream.write_all(data.as_bytes())?;
428            debug!("saved into stream");
429            Saved { world: saved_world }
430        }
431        SaveOutput::Drop => {
432            debug!("saved data dropped");
433            Saved { world: saved_world }
434        }
435        SaveOutput::Invalid => {
436            panic!("SaveOutput is invalid");
437        }
438    };
439
440    let result = Ok(saved);
441    event.after_save(world, &result);
442    result
443}
444
445// TODO: Documentation
446#[doc(hidden)]
447pub struct SaveCommand<E>(pub E);
448
449impl<E: SaveEvent> Command for SaveCommand<E> {
450    type Out = Result<(), SaveError>;
451
452    fn apply(self, world: &mut World) -> Result<(), SaveError> {
453        let saved = save_world(self.0, world)?;
454        world.trigger(saved);
455        Ok(())
456    }
457}
458
459#[cfg(test)]
460mod tests {
461    use std::fs::*;
462
463    use bevy::prelude::*;
464    use bevy_ecs::system::RunSystemOnce;
465
466    use super::*;
467
468    #[derive(Component, Default, Reflect)]
469    #[reflect(Component)]
470    #[require(Save)]
471    struct Foo;
472
473    fn app() -> App {
474        let mut app = App::new();
475        app.add_plugins(MinimalPlugins).register_type::<Foo>();
476        app
477    }
478
479    #[test]
480    fn test_save_into_file() {
481        #[derive(Resource)]
482        struct EventTriggered;
483
484        pub const PATH: &str = "test_save_into_file.ron";
485        let mut app = app();
486        app.add_observer(save_on_default_event);
487
488        app.add_observer(|_: On<Saved>, mut commands: Commands| {
489            commands.insert_resource(EventTriggered);
490        });
491
492        let _ = app.world_mut().run_system_once(|mut commands: Commands| {
493            commands.spawn((Foo, Save));
494            commands.trigger_save(SaveWorld::default_into_file(PATH));
495        });
496
497        let data = read_to_string(PATH).unwrap();
498        let world = app.world();
499        assert!(data.contains("Foo"));
500        assert!(world.contains_resource::<EventTriggered>());
501
502        remove_file(PATH).unwrap();
503    }
504
505    #[test]
506    fn test_save_into_stream() {
507        pub const PATH: &str = "test_save_to_stream.ron";
508
509        let mut app = app();
510        app.add_observer(save_on_default_event);
511
512        let _ = app.world_mut().run_system_once(|mut commands: Commands| {
513            commands.spawn((Foo, Save));
514            commands.trigger_save(SaveWorld::default_into_stream(File::create(PATH).unwrap()));
515        });
516
517        let data = read_to_string(PATH).unwrap();
518        assert!(data.contains("Foo"));
519
520        remove_file(PATH).unwrap();
521    }
522
523    #[test]
524    fn test_save_resource() {
525        pub const PATH: &str = "test_save_resource.ron";
526
527        #[derive(Resource, Default, Reflect)]
528        #[reflect(Resource)]
529        struct Bar;
530
531        let mut app = app();
532        app.register_type::<Bar>()
533            .add_observer(save_on_default_event);
534
535        let _ = app.world_mut().run_system_once(|mut commands: Commands| {
536            commands.insert_resource(Bar);
537            commands.trigger_save(
538                SaveWorld::default_into_stream(File::create(PATH).unwrap())
539                    .include_resource::<Bar>(),
540            );
541        });
542
543        app.update();
544
545        let data = read_to_string(PATH).unwrap();
546        assert!(data.contains("Bar"));
547
548        remove_file(PATH).unwrap();
549    }
550
551    #[test]
552    fn test_save_without_component() {
553        pub const PATH: &str = "test_save_without_component.ron";
554
555        #[derive(Component, Default, Reflect)]
556        #[reflect(Component)]
557        #[require(Save)]
558        struct Baz;
559
560        let mut app = app();
561        app.add_observer(save_on_default_event);
562
563        let _ = app.world_mut().run_system_once(|mut commands: Commands| {
564            commands.spawn((Foo, Baz, Save));
565            commands.trigger_save(SaveWorld::default_into_file(PATH).exclude_component::<Baz>());
566        });
567
568        let data = read_to_string(PATH).unwrap();
569        assert!(data.contains("Foo"));
570        assert!(!data.contains("Baz"));
571
572        remove_file(PATH).unwrap();
573    }
574
575    #[test]
576    fn test_map_component() {
577        pub const PATH: &str = "test_map_component.ron";
578
579        #[derive(Component, Default)]
580        struct Bar(#[allow(dead_code)] u32); // Not serializable
581
582        #[derive(Component, Default, Reflect)]
583        #[reflect(Component)]
584        struct Baz(u32); // Serializable
585
586        let mut app = app();
587        app.register_type::<Baz>()
588            .add_observer(save_on_default_event);
589
590        let entity = app
591            .world_mut()
592            .run_system_once(|mut commands: Commands| {
593                let entity = commands.spawn((Bar(12), Save)).id();
594                commands.trigger_save(
595                    SaveWorld::default_into_file(PATH).map_component::<Bar>(|Bar(i): &Bar| Baz(*i)),
596                );
597                entity
598            })
599            .unwrap();
600
601        let data = read_to_string(PATH).unwrap();
602        assert!(data.contains("Baz"));
603        assert!(data.contains("(12)"));
604        assert!(!data.contains("Bar"));
605        assert!(app.world().entity(entity).contains::<Bar>());
606        assert!(!app.world().entity(entity).contains::<Baz>());
607
608        remove_file(PATH).unwrap();
609    }
610}