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

use crate::prelude::{Children, Parent, PreviousParent};
use bevy_ecs::{Component, DynamicBundle, Entity, WorldBuilder};

pub struct WorldChildBuilder<'a, 'b> {
    world_builder: &'b mut WorldBuilder<'a>,
    parent_entities: Vec<Entity>,
}

impl<'a, 'b> WorldChildBuilder<'a, 'b> {
    pub fn spawn(&mut self, components: impl DynamicBundle + Send + Sync + 'static) -> &mut Self {
        let parent_entity = self
            .parent_entities
            .last()
            .cloned()
            .expect("There should always be a parent at this point.");
        self.world_builder
            .spawn(components)
            .with_bundle((Parent(parent_entity), PreviousParent(Some(parent_entity))));
        let entity = self.world_builder.current_entity.unwrap();
        {
            let world = &mut self.world_builder.world;
            let mut added = false;
            if let Ok(mut children) = world.get_mut::<Children>(parent_entity) {
                children.push(entity);
                added = true;
            }

            // NOTE: ideally this is just an else statement, but currently that _incorrectly_ fails borrow-checking
            if !added {
                world
                    .insert_one(parent_entity, Children(smallvec::smallvec![entity]))
                    .unwrap();
            }
        }
        self
    }

    pub fn with_bundle(
        &mut self,
        components: impl DynamicBundle + Send + Sync + 'static,
    ) -> &mut Self {
        self.world_builder.with_bundle(components);
        self
    }

    pub fn with(&mut self, component: impl Component) -> &mut Self {
        self.world_builder.with(component);
        self
    }
}

pub trait BuildWorldChildren {
    fn with_children(&mut self, spawn_children: impl FnMut(&mut WorldChildBuilder)) -> &mut Self;
}

impl<'a> BuildWorldChildren for WorldBuilder<'a> {
    fn with_children(
        &mut self,
        mut spawn_children: impl FnMut(&mut WorldChildBuilder),
    ) -> &mut Self {
        {
            let current_entity = self.current_entity.expect("Cannot add children because the 'current entity' is not set. You should spawn an entity first.");
            let mut builder = WorldChildBuilder {
                world_builder: self,
                parent_entities: vec![current_entity],
            };

            spawn_children(&mut builder);
        }
        self
    }
}

impl<'a, 'b> BuildWorldChildren for WorldChildBuilder<'a, 'b> {
    fn with_children(
        &mut self,
        mut spawn_children: impl FnMut(&mut WorldChildBuilder),
    ) -> &mut Self {
        let current_entity = self
            .world_builder
            .current_entity
            .expect("Cannot add children without a parent. Try creating an entity first.");
        self.parent_entities.push(current_entity);
        self.world_builder.current_entity = None;

        spawn_children(self);

        self.world_builder.current_entity = self.parent_entities.pop();
        self
    }
}