transform|ui: fix transform update lag

crates/bevy_app/src/app_builder.rs

@@ -164,6 +164,15 @@ impl AppBuilder {
         self
     }
 
+    pub fn add_system_to_stage_front(
+        &mut self,
+        stage_name: &'static str,
+        system: Box<dyn System>,
+    ) -> &mut Self {
+        self.app.schedule.add_system_to_stage_front(stage_name, system);
+        self
+    }
+
     pub fn add_systems_to_stage(
         &mut self,
         stage_name: &'static str,

crates/bevy_ecs/src/schedule/schedule.rs

@@ -98,6 +98,30 @@ impl Schedule {
         self
     }
 
+    pub fn add_system_to_stage_front(
+        &mut self,
+        stage_name: impl Into<Cow<'static, str>>,
+        system: Box<dyn System>,
+    ) -> &mut Self {
+        let stage_name = stage_name.into();
+        let systems = self
+            .stages
+            .get_mut(&stage_name)
+            .unwrap_or_else(|| panic!("Stage does not exist: {}", stage_name));
+        if self.system_ids.contains(&system.id()) {
+            panic!(
+                "System with id {:?} ({}) already exists",
+                system.id(),
+                system.name()
+            );
+        }
+        self.system_ids.insert(system.id());
+        systems.insert(0, Arc::new(Mutex::new(system)));
+
+        self.generation += 1;
+        self
+    }
+
     pub fn run(&mut self, world: &mut World, resources: &mut Resources) {
         for stage_name in self.stage_order.iter() {
             if let Some(stage_systems) = self.stages.get_mut(stage_name) {

crates/bevy_ecs/src/system/commands.rs

@@ -307,13 +307,13 @@ impl Commands {
         commands.current_entity
     }
 
-    pub fn for_current_entity(&mut self, mut func: impl FnMut(Entity)) -> &mut Self {
+    pub fn for_current_entity(&mut self, mut f: impl FnMut(Entity)) -> &mut Self {
         {
             let commands = self.commands.lock().unwrap();
             let current_entity = commands
                 .current_entity
                 .expect("The 'current entity' is not set. You should spawn an entity first.");
-            func(current_entity);
+            f(current_entity);
         }
         self
     }

crates/bevy_transform/src/hierarchy/child_builder.rs

@@ -16,7 +16,7 @@ impl WorldWriter for InsertChildren {
                     *child,
                     (
                         Parent(self.parent),
-                        PreviousParent(None),
+                        PreviousParent(Some(self.parent)),
                         LocalTransform::default(),
                     ),
                 )
@@ -57,7 +57,7 @@ impl WorldWriter for PushChildren {
                     *child,
                     (
                         Parent(self.parent),
-                        PreviousParent(None),
+                        PreviousParent(Some(self.parent)),
                         LocalTransform::default(),
                     ),
                 )
@@ -119,7 +119,7 @@ impl<'a> ChildBuilder<'a> {
 }
 
 pub trait BuildChildren {
-    fn with_children(&mut self, parent: impl FnMut(&mut ChildBuilder)) -> &mut Self;
+    fn with_children(&mut self, f: impl FnMut(&mut ChildBuilder)) -> &mut Self;
     fn push_children(&mut self, parent: Entity, children: &[Entity]) -> &mut Self;
     fn insert_children(&mut self, parent: Entity, index: usize, children: &[Entity]) -> &mut Self;
 }
@@ -252,11 +252,11 @@ mod tests {
 
         assert_eq!(
             *world.get::<PreviousParent>(child1).unwrap(),
-            PreviousParent(None)
+            PreviousParent(Some(parent))
         );
         assert_eq!(
             *world.get::<PreviousParent>(child2).unwrap(),
-            PreviousParent(None)
+            PreviousParent(Some(parent))
         );
 
         assert!(world.get::<LocalTransform>(child1).is_ok());
@@ -291,11 +291,11 @@ mod tests {
 
         assert_eq!(
             *world.get::<PreviousParent>(child1).unwrap(),
-            PreviousParent(None)
+            PreviousParent(Some(parent))
         );
         assert_eq!(
             *world.get::<PreviousParent>(child2).unwrap(),
-            PreviousParent(None)
+            PreviousParent(Some(parent))
         );
 
         assert!(world.get::<LocalTransform>(child1).is_ok());
@@ -313,11 +313,11 @@ mod tests {
         assert_eq!(*world.get::<Parent>(child4).unwrap(), Parent(parent));
         assert_eq!(
             *world.get::<PreviousParent>(child3).unwrap(),
-            PreviousParent(None)
+            PreviousParent(Some(parent))
         );
         assert_eq!(
             *world.get::<PreviousParent>(child4).unwrap(),
-            PreviousParent(None)
+            PreviousParent(Some(parent))
         );
 
         assert!(world.get::<LocalTransform>(child3).is_ok());

crates/bevy_transform/src/hierarchy/hierarchy_maintenance_system.rs

@@ -36,7 +36,7 @@ pub fn parent_update_system(
     }
 
     // Tracks all newly created `Children` Components this frame.
-    let mut children_additions = HashMap::<Entity, SmallVec<[Entity; 8]>>::with_capacity(16);
+    let mut children_additions = HashMap::<Entity, SmallVec<[Entity; 8]>>::new();
 
     // Entities with a changed Parent (that also have a PreviousParent, even if None)
     for (entity, parent, mut previous_parent) in &mut changed_parent_query.iter() {
@@ -108,7 +108,7 @@ pub fn hierarchy_maintenance_systems() -> Vec<Box<dyn System>> {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::transform_systems;
+    use crate::{hierarchy::BuildChildren, transform_systems};
     use bevy_ecs::{Resources, Schedule, World};
 
     #[test]
@@ -123,26 +123,21 @@ mod test {
         }
 
         // Add parent entities
-        let parent = world.spawn((Translation::new(1.0, 0.0, 0.0), Transform::identity()));
+        let mut commands = Commands::default();
-        let children = world
+        let mut parent = None;
-            .spawn_batch(vec![
+        let mut children = Vec::new();
-                (
+        commands
-                    Translation::new(0.0, 2.0, 0.0),
+            .spawn((Translation::new(1.0, 0.0, 0.0), Transform::identity()))
-                    LocalTransform::identity(),
+            .for_current_entity(|entity| parent = Some(entity))
-                    Transform::identity(),
+            .with_children(|parent| {
-                    Parent(parent),
+                parent
-                ),
+                    .spawn((Translation::new(0.0, 2.0, 0.0), Transform::identity()))
-                (
+                    .for_current_entity(|entity| children.push(entity))
-                    Translation::new(0.0, 0.0, 3.0),
+                    .spawn((Translation::new(0.0, 0.0, 3.0), Transform::identity()))
-                    LocalTransform::identity(),
+                    .for_current_entity(|entity| children.push(entity));
-                    Transform::identity(),
+            });
-                    Parent(parent),
+        let parent = parent.unwrap();
-                ),
+        commands.apply(&mut world, &mut resources);
-            ])
-            .collect::<Vec<Entity>>();
-
-        schedule.run(&mut world, &mut resources);
-        // TODO: this should be in sync after one run
         schedule.run(&mut world, &mut resources);
 
         assert_eq!(

crates/bevy_transform/src/hierarchy/world_child_builder.rs

@@ -25,7 +25,7 @@ impl<'a, 'b> WorldChildBuilder<'a, 'b> {
             .spawn_as_entity(entity, components)
             .with_bundle((
                 Parent(parent_entity),
-                PreviousParent(None),
+                PreviousParent(Some(parent_entity)),
                 LocalTransform::default(),
             ));
         {

crates/bevy_transform/src/lib.rs

@@ -39,6 +39,6 @@ impl AppPlugin for TransformPlugin {
             .register_component::<NonUniformScale>()
             // add transform systems to startup so the first update is "correct"
             .add_startup_systems(transform_systems())
-            .add_systems(transform_systems());
+            .add_systems_to_stage(stage::POST_UPDATE, transform_systems());
     }
 }

crates/bevy_transform/src/transform_propagate_system.rs

@@ -62,7 +62,7 @@ fn propagate_recursive(
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::transform_systems;
+    use crate::{hierarchy::BuildChildren, transform_systems};
     use bevy_ecs::{Resources, Schedule, World};
     use bevy_math::{Mat4, Vec3};
 
@@ -95,9 +95,9 @@ mod test {
                 ),
             ])
             .collect::<Vec<Entity>>();
-
+        // we need to run the schedule three times because components need to be filled in
-        // TODO: ideally we dont need three runs to keep transforms in sync.
+        // to resolve this problem in code, just add the correct components, or use Commands
-        // command buffers should be flushed in the appropriate places
+        // which adds all of the components needed with the correct state (see next test)
         schedule.run(&mut world, &mut resources);
         schedule.run(&mut world, &mut resources);
         schedule.run(&mut world, &mut resources);
@@ -114,4 +114,43 @@ mod test {
                 * Mat4::from_translation(Vec3::new(0.0, 0.0, 3.0))
         );
     }
+
+    #[test]
+    fn did_propagate_command_buffer() {
+        let mut world = World::default();
+        let mut resources = Resources::default();
+
+        let mut schedule = Schedule::default();
+        schedule.add_stage("update");
+        for system in transform_systems() {
+            schedule.add_system_to_stage("update", system);
+        }
+
+        // Root entity
+        let mut commands = Commands::default();
+        let mut children = Vec::new();
+        commands
+            .spawn((Translation::new(1.0, 0.0, 0.0), Transform::identity()))
+            .with_children(|parent| {
+                parent
+                    .spawn((Translation::new(0.0, 2.0, 0.0), Transform::identity()))
+                    .for_current_entity(|entity| children.push(entity))
+                    .spawn((Translation::new(0.0, 0.0, 3.0), Transform::identity()))
+                    .for_current_entity(|entity| children.push(entity));
+            });
+        commands.apply(&mut world, &mut resources);
+        schedule.run(&mut world, &mut resources);
+
+        assert_eq!(
+            world.get::<Transform>(children[0]).unwrap().value,
+            Mat4::from_translation(Vec3::new(1.0, 0.0, 0.0))
+                * Mat4::from_translation(Vec3::new(0.0, 2.0, 0.0))
+        );
+
+        assert_eq!(
+            world.get::<Transform>(children[1]).unwrap().value,
+            Mat4::from_translation(Vec3::new(1.0, 0.0, 0.0))
+                * Mat4::from_translation(Vec3::new(0.0, 0.0, 3.0))
+        );
+    }
 }

crates/bevy_ui/src/lib.rs

@@ -32,7 +32,8 @@ pub struct UiPlugin;
 impl AppPlugin for UiPlugin {
     fn build(&self, app: &mut AppBuilder) {
         app.add_system_to_stage(stage::PRE_UPDATE, ui_focus_system.system())
-            .add_system_to_stage(stage::POST_UPDATE, ui_update_system.system())
+            // must run before transform update systems
+            .add_system_to_stage_front(stage::POST_UPDATE, ui_update_system.system())
             .add_system_to_stage(stage::POST_UPDATE, widget::text_system.system())
             .add_system_to_stage(bevy_render::stage::DRAW, widget::draw_text_system.system());