From 4a6b686832a9ea19bc6738f4019bb58d9d052ee3 Mon Sep 17 00:00:00 2001
From: andriyDev <andriydzikh@gmail.com>
Date: Tue, 19 Nov 2024 08:31:00 -0800
Subject: [PATCH] Fix runtime required components not registering correctly
 (#16436)

# Objective

- Fixes #16406
- Fixes an issue where registering a "deeper" required component, then a
"shallower" required component, would result in the wrong required
constructor being used for the root component.

## Solution

- Make `register_required_components` add any "parent" of a component as
`required_by` to the new "child".
- Assign the depth of the `requiree` plus 1 as the depth of a new
runtime required component.

## Testing

- Added two new tests.
---
 crates/bevy_ecs/src/component.rs | 11 +++++--
 crates/bevy_ecs/src/lib.rs       | 54 ++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/crates/bevy_ecs/src/component.rs b/crates/bevy_ecs/src/component.rs
index 7c7df1b90a..d3b72f5a3a 100644
--- a/crates/bevy_ecs/src/component.rs
+++ b/crates/bevy_ecs/src/component.rs
@@ -1065,6 +1065,10 @@ impl Components {
 
         // Propagate the new required components up the chain to all components that require the requiree.
         if let Some(required_by) = self.get_required_by(requiree).cloned() {
+            // `required` is now required by anything that `requiree` was required by.
+            self.get_required_by_mut(required)
+                .unwrap()
+                .extend(required_by.iter().copied());
             for &required_by_id in required_by.iter() {
                 // SAFETY: The component is in the list of required components, so it must exist already.
                 let required_components = unsafe {
@@ -1072,9 +1076,10 @@ impl Components {
                         .debug_checked_unwrap()
                 };
 
-                // Register the original required component for the requiree.
-                // The inheritance depth is `1` since this is a component required by the original requiree.
-                required_components.register_by_id(required, constructor, 1);
+                // Register the original required component in the "parent" of the requiree.
+                // The inheritance depth is 1 deeper than the `requiree` wrt `required_by_id`.
+                let depth = required_components.0.get(&requiree).expect("requiree is required by required_by_id, so its required_components must include requiree").inheritance_depth;
+                required_components.register_by_id(required, constructor, depth + 1);
 
                 for (component_id, component) in inherited_requirements.iter() {
                     // Register the required component.
diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs
index f18bde71f7..4eab4d00be 100644
--- a/crates/bevy_ecs/src/lib.rs
+++ b/crates/bevy_ecs/src/lib.rs
@@ -2345,6 +2345,60 @@ mod tests {
         assert!(world.entity(id).get::<C>().is_some());
     }
 
+    #[test]
+    fn runtime_required_components_propagate_up_even_more() {
+        #[derive(Component)]
+        struct A;
+
+        #[derive(Component, Default)]
+        struct B;
+
+        #[derive(Component, Default)]
+        struct C;
+
+        #[derive(Component, Default)]
+        struct D;
+
+        let mut world = World::new();
+
+        world.register_required_components::<A, B>();
+        world.register_required_components::<B, C>();
+        world.register_required_components::<C, D>();
+
+        let id = world.spawn(A).id();
+
+        assert!(world.entity(id).get::<D>().is_some());
+    }
+
+    #[test]
+    fn runtime_required_components_deep_require_does_not_override_shallow_require() {
+        #[derive(Component)]
+        struct A;
+        #[derive(Component, Default)]
+        struct B;
+        #[derive(Component, Default)]
+        struct C;
+        #[derive(Component)]
+        struct Counter(i32);
+        #[derive(Component, Default)]
+        struct D;
+
+        let mut world = World::new();
+
+        world.register_required_components::<A, B>();
+        world.register_required_components::<B, C>();
+        world.register_required_components::<C, D>();
+        world.register_required_components_with::<D, Counter>(|| Counter(2));
+        // This should replace the require constructor in A since it is
+        // shallower.
+        world.register_required_components_with::<C, Counter>(|| Counter(1));
+
+        let id = world.spawn(A).id();
+
+        // The "shallower" of the two components is used.
+        assert_eq!(world.entity(id).get::<Counter>().unwrap().0, 1);
+    }
+
     #[test]
     fn runtime_required_components_existing_archetype() {
         #[derive(Component)]