Heritable trait based property inheritance systems

crates/bevy_hierarchy/src/inheritance.rs

@@ -0,0 +1,107 @@
+use crate::{Children, HierarchySystem, Parent};
+use bevy_app::prelude::*;
+use bevy_ecs::prelude::*;
+
+/// Defines a component that inherits state from it's ancestors.
+/// These types are typically publicly read-only, relying on a "source"
+/// companion component that is used to compute the new state of a by
+/// composing the source and it's parent.
+pub trait Heritable: Component + Copy {
+    /// The source component where the base state is sourced from.
+    type Source: Component;
+    /// Updates the base state of a root-level component based on the companion
+    /// source component.
+    fn root(&mut self, source: &Self::Source);
+    /// Updates the a mid-level or leaf component in a hierarchy based on the
+    /// companion source component and the immediate parent of the entity.
+    fn inherit(&mut self, parent: &Self, source: &Self::Source);
+}
+
+/// Extension trate for adding inheritance based systems to an [`App`].
+pub trait HeritableAppExt {
+    /// Registers systems for propagating the inherited states.
+    /// See [`Heritable`] for more information about hierarchical inheritance.
+    fn register_heritable<T: Heritable>(&mut self) -> &mut Self;
+}
+
+impl HeritableAppExt for App {
+    fn register_heritable<T: Heritable>(&mut self) -> &mut Self {
+        // Adding these to startup ensures the first update is "correct"
+        self.add_startup_system_to_stage(
+            StartupStage::PostStartup,
+            inheritance_system::<T>
+                .label(HierarchySystem::InheritancePropagation)
+                .after(HierarchySystem::ParentUpdate),
+        )
+        .add_system_to_stage(
+            CoreStage::PostUpdate,
+            inheritance_system::<T>
+                .label(HierarchySystem::InheritancePropagation)
+                .after(HierarchySystem::ParentUpdate),
+        )
+    }
+}
+
+/// Update children in a hierarchy based on the properties of their parents.
+pub fn inheritance_system<T: Heritable>(
+    mut root_query: Query<(Entity, Option<&Children>, &T::Source, &mut T), Without<Parent>>,
+    mut source_query: Query<(&T::Source, &mut T), With<Parent>>,
+    changed_query: Query<Entity, Changed<T::Source>>,
+    children_query: Query<Option<&Children>, (With<Parent>, With<T>)>,
+) {
+    for (entity, children, source, mut root_component) in root_query.iter_mut() {
+        let mut changed = false;
+        if changed_query.get(entity).is_ok() {
+            root_component.root(source);
+            changed = true;
+        }
+
+        if let Some(children) = children {
+            for child in children.iter() {
+                propagate_recursive(
+                    &*root_component,
+                    &changed_query,
+                    &mut source_query,
+                    &children_query,
+                    *child,
+                    changed,
+                );
+            }
+        }
+    }
+}
+
+fn propagate_recursive<T: Heritable>(
+    parent: &T,
+    changed_query: &Query<Entity, Changed<T::Source>>,
+    source_query: &mut Query<(&T::Source, &mut T), With<Parent>>,
+    children_query: &Query<Option<&Children>, (With<Parent>, With<T>)>,
+    entity: Entity,
+    mut changed: bool,
+) {
+    changed |= changed_query.get(entity).is_ok();
+
+    let component = {
+        if let Ok((source, mut component)) = source_query.get_mut(entity) {
+            if changed {
+                component.inherit(parent, source);
+            }
+            *component
+        } else {
+            return;
+        }
+    };
+
+    if let Ok(Some(children)) = children_query.get(entity) {
+        for child in children.iter() {
+            propagate_recursive(
+                &component,
+                changed_query,
+                source_query,
+                children_query,
+                *child,
+                changed,
+            );
+        }
+    }
+}

crates/bevy_hierarchy/src/lib.rs

@@ -13,13 +13,22 @@ pub use hierarchy::*;
 mod child_builder;
 pub use child_builder::*;
 
+/// Traits, systems, and plugins for hierarchical inheritance.
+pub mod inheritance;
+
 mod systems;
 pub use systems::*;
 
 #[doc(hidden)]
 pub mod prelude {
     #[doc(hidden)]
-    pub use crate::{child_builder::*, components::*, hierarchy::*, HierarchyPlugin};
+    pub use crate::{
+        child_builder::*,
+        components::*,
+        hierarchy::*,
+        inheritance::{Heritable, HeritableAppExt},
+        HierarchyPlugin,
+    };
 }
 
 use bevy_app::prelude::*;
@@ -34,6 +43,8 @@ pub struct HierarchyPlugin;
 pub enum HierarchySystem {
     /// Updates [`Parent`] when changes in the hierarchy occur
     ParentUpdate,
+    /// Propagates inherited properties from [`Parent`] to [`Children`].
+    InheritancePropagation,
 }
 
 impl Plugin for HierarchyPlugin {

crates/bevy_transform/src/components/global_transform.rs

@@ -1,5 +1,6 @@
 use super::Transform;
 use bevy_ecs::{component::Component, reflect::ReflectComponent};
+use bevy_hierarchy::inheritance::Heritable;
 use bevy_math::{const_vec3, Mat3, Mat4, Quat, Vec3};
 use bevy_reflect::Reflect;
 use std::ops::Mul;
@@ -19,7 +20,7 @@ use std::ops::Mul;
 /// [`GlobalTransform`] is the position of an entity relative to the reference frame.
 ///
 /// [`GlobalTransform`] is updated from [`Transform`] in the system
-/// [`transform_propagate_system`](crate::transform_propagate_system).
+/// [`inheritance_system`](bevy_hierarchy::inheritance::inheritance_system).
 ///
 /// This system runs in stage [`CoreStage::PostUpdate`](crate::CoreStage::PostUpdate). If you
 /// update the[`Transform`] of an entity in this stage or after, you will notice a 1 frame lag
@@ -277,3 +278,204 @@ impl Mul<Vec3> for GlobalTransform {
         self.mul_vec3(value)
     }
 }
+
+impl Heritable for GlobalTransform {
+    type Source = Transform;
+    fn root(&mut self, source: &Transform) {
+        *self = GlobalTransform::from(*source)
+    }
+
+    fn inherit(&mut self, parent: &GlobalTransform, source: &Transform) {
+        *self = parent.mul_transform(*source);
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use bevy_ecs::{
+        schedule::{Schedule, Stage, SystemStage},
+        system::{CommandQueue, Commands},
+        world::World,
+    };
+
+    use crate::components::{GlobalTransform, Transform};
+    use crate::TransformBundle;
+    use bevy_hierarchy::{
+        inheritance::inheritance_system, parent_update_system, BuildChildren, BuildWorldChildren,
+        Children, Parent,
+    };
+
+    #[test]
+    fn did_propagate() {
+        let mut world = World::default();
+
+        let mut update_stage = SystemStage::parallel();
+        update_stage.add_system(parent_update_system);
+        update_stage.add_system(inheritance_system::<GlobalTransform>);
+
+        let mut schedule = Schedule::default();
+        schedule.add_stage("update", update_stage);
+
+        // Root entity
+        world
+            .spawn()
+            .insert_bundle(TransformBundle::from(Transform::from_xyz(1.0, 0.0, 0.0)));
+
+        let mut children = Vec::new();
+        world
+            .spawn()
+            .insert_bundle(TransformBundle::from(Transform::from_xyz(1.0, 0.0, 0.0)))
+            .with_children(|parent| {
+                children.push(
+                    parent
+                        .spawn_bundle(TransformBundle::from(Transform::from_xyz(0.0, 2.0, 0.)))
+                        .id(),
+                );
+                children.push(
+                    parent
+                        .spawn_bundle(TransformBundle::from(Transform::from_xyz(0.0, 0.0, 3.)))
+                        .id(),
+                );
+            });
+        schedule.run(&mut world);
+
+        assert_eq!(
+            *world.get::<GlobalTransform>(children[0]).unwrap(),
+            GlobalTransform::from_xyz(1.0, 0.0, 0.0) * Transform::from_xyz(0.0, 2.0, 0.0)
+        );
+
+        assert_eq!(
+            *world.get::<GlobalTransform>(children[1]).unwrap(),
+            GlobalTransform::from_xyz(1.0, 0.0, 0.0) * Transform::from_xyz(0.0, 0.0, 3.0)
+        );
+    }
+
+    #[test]
+    fn did_propagate_command_buffer() {
+        let mut world = World::default();
+
+        let mut update_stage = SystemStage::parallel();
+        update_stage.add_system(parent_update_system);
+        update_stage.add_system(inheritance_system::<GlobalTransform>);
+
+        let mut schedule = Schedule::default();
+        schedule.add_stage("update", update_stage);
+
+        // Root entity
+        let mut queue = CommandQueue::default();
+        let mut commands = Commands::new(&mut queue, &world);
+        let mut children = Vec::new();
+        commands
+            .spawn_bundle(TransformBundle::from(Transform::from_xyz(1.0, 0.0, 0.0)))
+            .with_children(|parent| {
+                children.push(
+                    parent
+                        .spawn_bundle(TransformBundle::from(Transform::from_xyz(0.0, 2.0, 0.0)))
+                        .id(),
+                );
+                children.push(
+                    parent
+                        .spawn_bundle(TransformBundle::from(Transform::from_xyz(0.0, 0.0, 3.0)))
+                        .id(),
+                );
+            });
+        queue.apply(&mut world);
+        schedule.run(&mut world);
+
+        assert_eq!(
+            *world.get::<GlobalTransform>(children[0]).unwrap(),
+            GlobalTransform::from_xyz(1.0, 0.0, 0.0) * Transform::from_xyz(0.0, 2.0, 0.0)
+        );
+
+        assert_eq!(
+            *world.get::<GlobalTransform>(children[1]).unwrap(),
+            GlobalTransform::from_xyz(1.0, 0.0, 0.0) * Transform::from_xyz(0.0, 0.0, 3.0)
+        );
+    }
+
+    #[test]
+    fn correct_children() {
+        let mut world = World::default();
+
+        let mut update_stage = SystemStage::parallel();
+        update_stage.add_system(parent_update_system);
+        update_stage.add_system(inheritance_system::<GlobalTransform>);
+
+        let mut schedule = Schedule::default();
+        schedule.add_stage("update", update_stage);
+
+        // Add parent entities
+        let mut command_queue = CommandQueue::default();
+        let mut commands = Commands::new(&mut command_queue, &world);
+        let mut children = Vec::new();
+        let parent = commands
+            .spawn()
+            .insert(Transform::from_xyz(1.0, 0.0, 0.0))
+            .id();
+        commands.entity(parent).with_children(|parent| {
+            children.push(
+                parent
+                    .spawn()
+                    .insert(Transform::from_xyz(0.0, 2.0, 0.0))
+                    .id(),
+            );
+            children.push(
+                parent
+                    .spawn()
+                    .insert(Transform::from_xyz(0.0, 3.0, 0.0))
+                    .id(),
+            );
+        });
+        command_queue.apply(&mut world);
+        schedule.run(&mut world);
+
+        assert_eq!(
+            world
+                .get::<Children>(parent)
+                .unwrap()
+                .iter()
+                .cloned()
+                .collect::<Vec<_>>(),
+            children,
+        );
+
+        // Parent `e1` to `e2`.
+        (*world.get_mut::<Parent>(children[0]).unwrap()).0 = children[1];
+
+        schedule.run(&mut world);
+
+        assert_eq!(
+            world
+                .get::<Children>(parent)
+                .unwrap()
+                .iter()
+                .cloned()
+                .collect::<Vec<_>>(),
+            vec![children[1]]
+        );
+
+        assert_eq!(
+            world
+                .get::<Children>(children[1])
+                .unwrap()
+                .iter()
+                .cloned()
+                .collect::<Vec<_>>(),
+            vec![children[0]]
+        );
+
+        assert!(world.despawn(children[0]));
+
+        schedule.run(&mut world);
+
+        assert_eq!(
+            world
+                .get::<Children>(parent)
+                .unwrap()
+                .iter()
+                .cloned()
+                .collect::<Vec<_>>(),
+            vec![children[1]]
+        );
+    }
+}

crates/bevy_transform/src/components/transform.rs

@@ -20,7 +20,7 @@ use std::ops::Mul;
 /// [`GlobalTransform`] is the position of an entity relative to the reference frame.
 ///
 /// [`GlobalTransform`] is updated from [`Transform`] in the system
-/// [`transform_propagate_system`](crate::transform_propagate_system).
+/// [`inheritance_system`](bevy_hierarchy::inheritance::inheritance_system).
 ///
 /// This system runs in stage [`CoreStage::PostUpdate`](crate::CoreStage::PostUpdate). If you
 /// update the[`Transform`] of an entity in this stage or after, you will notice a 1 frame lag