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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183

use std::ops::{Deref, DerefMut};

use bevy::{
    ecs::system::{Command, EntityCommands},
    prelude::*,
};
use smallvec::SmallVec;

use crate::{prelude::*, render::ShapePipelineType};

/// Command that pushes children to the end of the entity's [`Children`].
///
/// Duplicated here from [`bevy::prelude::PushChildren`] in order to access private internals.
#[derive(Debug)]
pub struct PushChildren {
    parent: Entity,
    children: SmallVec<[Entity; 8]>,
}

impl Command for PushChildren {
    fn write(self, world: &mut World) {
        world.entity_mut(self.parent).push_children(&self.children);
    }
}

/// [`EntityCommands`] that also stores [`ShapeConfig`] for easier spawning of child shapes.
pub struct ShapeEntityCommands<'w, 's, 'a> {
    pub commands: EntityCommands<'w, 's, 'a>,
    pub config: &'a ShapeConfig,
}

impl<'w, 's, 'a> ShapeEntityCommands<'w, 's, 'a> {
    /// Takes a closure which builds children for this entity using [`ShapeChildBuilder`].
    pub fn with_children(
        &mut self,
        spawn_children: impl FnOnce(&mut ShapeChildBuilder),
    ) -> &mut Self {
        let config = self.config.without_transform();
        let parent = self.id();
        let mut painter = ShapeChildBuilder {
            commands: self.commands(),
            push_children: PushChildren {
                children: SmallVec::default(),
                parent,
            },
            config,
        };

        spawn_children(&mut painter);
        let children = painter.push_children;
        self.commands().add(children);
        self
    }
}

impl<'w, 's, 'a> Deref for ShapeEntityCommands<'w, 's, 'a> {
    type Target = EntityCommands<'w, 's, 'a>;

    fn deref(&self) -> &Self::Target {
        &self.commands
    }
}

impl<'w, 's, 'a> DerefMut for ShapeEntityCommands<'w, 's, 'a> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.commands
    }
}

/// [`ChildBuilder`] that also exposes shape spawning methods from [`ShapeCommands`].
pub struct ShapeChildBuilder<'w, 's, 'a> {
    commands: &'a mut Commands<'w, 's>,
    config: ShapeConfig,
    push_children: PushChildren,
}

impl<'w, 's, 'a> ShapeChildBuilder<'w, 's, 'a> {
    /// Spawns an entity with the given bundle and inserts it into the parent entity's [`Children`].
    /// Also adds [`Parent`] component to the created entity.
    pub fn spawn(&mut self, bundle: impl Bundle) -> EntityCommands<'w, 's, '_> {
        let e = self.commands.spawn(bundle);
        self.push_children.children.push(e.id());
        e
    }

    /// Spawns an [`Entity`] with no components and inserts it into the parent entity's [`Children`].
    /// Also adds [`Parent`] component to the created entity.
    pub fn spawn_empty(&mut self) -> EntityCommands<'w, 's, '_> {
        let e = self.commands.spawn_empty();
        self.push_children.children.push(e.id());
        e
    }

    /// Returns the parent entity of this [`ChildBuilder`].
    pub fn parent_entity(&self) -> Entity {
        self.push_children.parent
    }

    /// Adds a command to be executed, like [`Commands::add`].
    pub fn add_command<C: Command + 'static>(&mut self, command: C) -> &mut Self {
        self.commands.add(command);
        self
    }
}

impl<'w, 's, 'a> ShapeSpawner<'w, 's> for ShapeChildBuilder<'w, 's, 'a> {
    fn spawn_shape(&mut self, bundle: impl Bundle) -> ShapeEntityCommands<'w, 's, '_> {
        let Self {
            commands, config, ..
        } = self;
        let mut e = commands.spawn(bundle);
        self.push_children.children.push(e.id());
        if let Some(layers) = config.render_layers {
            e.insert(layers);
        }
        if let ShapePipelineType::Shape3d = config.pipeline {
            e.insert(Shape3d);
        }

        ShapeEntityCommands {
            commands: e,
            config,
        }
    }

    fn config(&self) -> &ShapeConfig {
        &self.config
    }

    fn set_config(&mut self, config: ShapeConfig) {
        self.config = config;
    }
}

impl<'w, 's, 'a> Deref for ShapeChildBuilder<'w, 's, 'a> {
    type Target = ShapeConfig;

    fn deref(&self) -> &Self::Target {
        &self.config
    }
}

impl<'w, 's, 'a> DerefMut for ShapeChildBuilder<'w, 's, 'a> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.config
    }
}

/// Extension trait for [`EntityCommands`] to allow injection of [`ShapeConfig`].
///
/// Useful when parenting shapes under a non-shape entity.
pub trait BuildShapeChildren {
    /// Similar to [`ShapeEntityCommands::with_children`] except is available on non-shape entities, takes in config to pass along to the [`ShapeChildBuilder`]
    fn with_shape_children(
        &mut self,
        config: &ShapeConfig,
        f: impl FnOnce(&mut ShapeChildBuilder),
    ) -> &mut Self;
}

impl<'w, 's, 'a> BuildShapeChildren for EntityCommands<'w, 's, 'a> {
    fn with_shape_children(
        &mut self,
        config: &ShapeConfig,
        spawn_children: impl FnOnce(&mut ShapeChildBuilder),
    ) -> &mut Self {
        let config = config.without_transform();
        let parent = self.id();
        let mut painter = ShapeChildBuilder {
            commands: self.commands(),
            push_children: PushChildren {
                children: SmallVec::default(),
                parent,
            },
            config,
        };

        spawn_children(&mut painter);
        let children = painter.push_children;
        self.commands().add(children);
        self
    }
}