From 33abd3e7f40a3a6d0788c8f669d76b035dadf4c5 Mon Sep 17 00:00:00 2001
From: UkoeHB <37489173+UkoeHB@users.noreply.github.com>
Date: Mon, 11 Nov 2024 15:59:56 -0600
Subject: [PATCH] Fix panic in UiSurface if Node is removed and re-added
 (#16288)

# Objective

- Fix bug where `UiSurface::set_camera_children` (and
`UiSurface::update_children` sometimes) will panic if you remove and add
a `Node` component in a single tick. This is more likely to happen now
because of `remove_with_requires`.

## Solution

- Filter out entities with `Node` when cleaning up entities from
`RemovedComponents<Node>`.

## Testing

- Not tested (rust compiler refused to cooperate when I tried to patch
this into my project), correct by inspection.
---
 crates/bevy_ui/src/layout/mod.rs | 39 +++++++++++++++++++++++++++++++-
 1 file changed, 38 insertions(+), 1 deletion(-)

diff --git a/crates/bevy_ui/src/layout/mod.rs b/crates/bevy_ui/src/layout/mod.rs
index c4a5d63880..aaa3f47050 100644
--- a/crates/bevy_ui/src/layout/mod.rs
+++ b/crates/bevy_ui/src/layout/mod.rs
@@ -278,7 +278,12 @@ with UI components as a child of an entity without UI components, your UI layout
 
     let text_buffers = &mut buffer_query;
     // clean up removed nodes after syncing children to avoid potential panic (invalid SlotMap key used)
-    ui_surface.remove_entities(removed_components.removed_nodes.read());
+    ui_surface.remove_entities(
+        removed_components
+            .removed_nodes
+            .read()
+            .filter(|entity| !node_query.contains(*entity)),
+    );
 
     // Re-sync changed children: avoid layout glitches caused by removed nodes that are still set as a child of another node
     computed_node_query.iter().for_each(|(entity, _)| {
@@ -771,6 +776,38 @@ mod tests {
         assert!(ui_surface.entity_to_taffy.is_empty());
     }
 
+    /// bugfix test, see [#16288](https://github.com/bevyengine/bevy/pull/16288)
+    #[test]
+    fn node_removal_and_reinsert_should_work() {
+        let (mut world, mut ui_schedule) = setup_ui_test_world();
+
+        ui_schedule.run(&mut world);
+
+        // no UI entities in world, none in UiSurface
+        let ui_surface = world.resource::<UiSurface>();
+        assert!(ui_surface.entity_to_taffy.is_empty());
+
+        let ui_entity = world.spawn(Node::default()).id();
+
+        // `ui_layout_system` should map `ui_entity` to a ui node in `UiSurface::entity_to_taffy`
+        ui_schedule.run(&mut world);
+
+        let ui_surface = world.resource::<UiSurface>();
+        assert!(ui_surface.entity_to_taffy.contains_key(&ui_entity));
+        assert_eq!(ui_surface.entity_to_taffy.len(), 1);
+
+        // remove and re-insert Node to trigger removal code in `ui_layout_system`
+        world.entity_mut(ui_entity).remove::<Node>();
+        world.entity_mut(ui_entity).insert(Node::default());
+
+        // `ui_layout_system` should still have `ui_entity`
+        ui_schedule.run(&mut world);
+
+        let ui_surface = world.resource::<UiSurface>();
+        assert!(ui_surface.entity_to_taffy.contains_key(&ui_entity));
+        assert_eq!(ui_surface.entity_to_taffy.len(), 1);
+    }
+
     /// regression test for >=0.13.1 root node layouts
     /// ensure root nodes act like they are absolutely positioned
     /// without explicitly declaring it.