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
142
143

use std::ops::RangeFrom;

use bevy::ecs::query::{QueryData, QueryFilter, WorldQuery};
use bevy::ecs::system::{EntityCommands, SystemParam};
use bevy::prelude::*;
use bevy::utils::HashMap;

use crate::entity_management::EntitiesToPopulate;

/// Wrapper for writing queries in populate systems.
#[derive(SystemParam)]
pub struct YoleckPopulate<'w, 's, Q: 'static + QueryData, F: 'static + QueryFilter = ()> {
    entities_to_populate: Res<'w, EntitiesToPopulate>,
    query: Query<'w, 's, Q, F>,
    commands: Commands<'w, 's>,
}

impl<Q: 'static + QueryData, F: 'static + QueryFilter> YoleckPopulate<'_, '_, Q, F> {
    /// Iterate over the entities that need populating in order to add/update components using
    /// a Bevy command.
    pub fn populate(
        &mut self,
        mut dlg: impl FnMut(YoleckPopulateContext, EntityCommands, <Q as WorldQuery>::Item<'_>),
    ) {
        for (entity, populate_reason) in self.entities_to_populate.0.iter() {
            if let Ok(data) = self.query.get_mut(*entity) {
                let cmd = self.commands.entity(*entity);
                let context = YoleckPopulateContext {
                    reason: *populate_reason,
                };
                dlg(context, cmd, data);
            }
        }
    }
}

#[derive(Clone, Copy)]
pub(crate) enum PopulateReason {
    EditorInit,
    EditorUpdate,
    RealGame,
}

/// A context for [`YoleckPopulate::populate`].
pub struct YoleckPopulateContext {
    pub(crate) reason: PopulateReason,
}

impl YoleckPopulateContext {
    /// `true` if the entity is created in editor mode, `false` if created in playtest or actual game.
    pub fn is_in_editor(&self) -> bool {
        match self.reason {
            PopulateReason::EditorInit => true,
            PopulateReason::EditorUpdate => true,
            PopulateReason::RealGame => false,
        }
    }

    /// `true` if this is this is the first time the entity is populated, `false` if the entity was
    /// popultated before.
    pub fn is_first_time(&self) -> bool {
        match self.reason {
            PopulateReason::EditorInit => true,
            PopulateReason::EditorUpdate => false,
            PopulateReason::RealGame => true,
        }
    }
}

/// See [`YoleckMarking`].
#[derive(Debug, Component, PartialEq, Eq, Clone, Copy)]
pub struct YoleckSystemMarker(usize);

#[derive(Resource)]
struct MarkerGenerator(RangeFrom<usize>);

impl FromWorld for YoleckSystemMarker {
    fn from_world(world: &mut World) -> Self {
        let mut marker = world.get_resource_or_insert_with(|| MarkerGenerator(1..));
        YoleckSystemMarker(marker.0.next().unwrap())
    }
}

/// Use to mark child entities created from a specific system.
///
/// Using `despawn_descendants` is dangerous because it can delete entities created by other
/// populate systems in the same frame. Instead, a populate system that wants to despawn its own
/// child entities should do something like this:
///
/// ```no_run
/// # use bevy::prelude::*;
/// # use bevy_yoleck::prelude::*;
/// # use serde::{Deserialize, Serialize};
/// # #[derive(Default, Clone, PartialEq, Serialize, Deserialize, Component, YoleckComponent)]
/// # struct MyComponent;
/// fn populate_system(mut populate: YoleckPopulate<&MyComponent>, marking: YoleckMarking) {
///     populate.populate(|_ctx, mut cmd, my_component| {
///         marking.despawn_marked(&mut cmd);
///         cmd.with_children(|commands| {
///             let mut child = commands.spawn(marking.marker());
///             child.insert((
///                 // relevant Bevy components
///             ));
///         });
///     });
/// }
/// ```
#[derive(SystemParam)]
pub struct YoleckMarking<'w, 's> {
    designated_marker: Local<'s, YoleckSystemMarker>,
    children_query: Query<'w, 's, &'static Children>,
    marked_query: Query<'w, 's, (&'static Parent, &'static YoleckSystemMarker)>,
}

impl YoleckMarking<'_, '_> {
    /// Create a marker unique to this system.
    pub fn marker(&self) -> YoleckSystemMarker {
        *self.designated_marker
    }

    /// Despawn (with `despawn_recursive`) all the entities marked by the current system, that are
    /// descendants of the entity edited by the supplied `cmd`.
    pub fn despawn_marked(&self, cmd: &mut EntityCommands) {
        let mut marked_children_map: HashMap<Entity, Vec<Entity>> = Default::default();
        for child in self.children_query.iter_descendants(cmd.id()) {
            let Ok((parent, marker)) = self.marked_query.get(child) else {
                continue;
            };
            if *marker == *self.designated_marker {
                marked_children_map
                    .entry(parent.get())
                    .or_default()
                    .push(child);
            }
        }
        for (parent, children) in marked_children_map {
            cmd.commands().entity(parent).remove_children(&children);
            for child in children {
                cmd.commands().entity(child).despawn_recursive();
            }
        }
    }
}