Fix Taffy viewport node leaks

[ickshonpe]

Objective

For most UI node entities there's a 1-to-1 mapping from the entity to its associated Taffy node. Root UI nodes are an exception though, their corresponding Taffy node in the Taffy tree is also given a parent that represents the viewport. These viewport Taffy nodes are not removed when a root UI node is despawned.

Parenting of an existing root UI node with an associated viewport Taffy node also results in the leak of the viewport node.

These tests fail if added to the layout module's tests on the main branch:

    #[test]
    fn no_viewport_node_leak_on_root_despawned() {
        let (mut world, mut ui_schedule) = setup_ui_test_world();

        let ui_root_entity = world.spawn(Node::default()).id();

        // The UI schedule synchronizes Bevy UI's internal `TaffyTree` with the
        // main world's tree of `Node` entities.
        ui_schedule.run(&mut world);

        // Two taffy nodes are added to the internal `TaffyTree` for each root UI entity.
        // An implicit taffy node representing the viewport and a taffy node corresponding to the
        // root UI entity which is parented to the viewport taffy node.
        assert_eq!(
            world.resource_mut::<UiSurface>().taffy.total_node_count(),
            2
        );

        world.despawn(ui_root_entity);

        // The UI schedule removes both the taffy node corresponding to `ui_root_entity` and its
        // parent viewport node.
        ui_schedule.run(&mut world);

        // Both taffy nodes should now be removed from the internal `TaffyTree`
        assert_eq!(
            world.resource_mut::<UiSurface>().taffy.total_node_count(),
            0
        );
    }

    #[test]
    fn no_viewport_node_leak_on_parented_root() {
        let (mut world, mut ui_schedule) = setup_ui_test_world();

        let ui_root_entity_1 = world.spawn(Node::default()).id();
        let ui_root_entity_2 = world.spawn(Node::default()).id();

        ui_schedule.run(&mut world);

        // There are two UI root entities. Each root taffy node is given it's own viewport node parent,
        // so a total of four taffy nodes are added to the `TaffyTree` by the UI schedule.
        assert_eq!(
            world.resource_mut::<UiSurface>().taffy.total_node_count(),
            4
        );

        // Parent `ui_root_entity_2` onto `ui_root_entity_1` so now only `ui_root_entity_1` is a
        // UI root entity.
        world
            .entity_mut(ui_root_entity_1)
            .add_child(ui_root_entity_2);

        // Now there is only one root node so the second viewport node is removed by
        // the UI schedule.
        ui_schedule.run(&mut world);

        // There is only one viewport node now, so the `TaffyTree` contains 3 nodes in total.
        assert_eq!(
            world.resource_mut::<UiSurface>().taffy.total_node_count(),
            3
        );
    }

Fixes #17594

Solution

Change the UiSurface::entity_to_taffy to map to LayoutNodes. A LayoutNode has a viewport_id: Option<taffy::NodeId> field which is the id of the corresponding implicit "viewport" node if the node is a root UI node, otherwise it is None. When removing or parenting nodes this field is checked and the implicit viewport node is removed if present.

Testing

There are two new tests in bevy_ui::layout::tests included with this PR:

• no_viewport_node_leak_on_root_despawned
• no_viewport_node_leak_on_parented_root

[mockersf]

adding it here means that it will allocate memory for each entry even though only few of them will be actually filled

[mockersf]

that's an additional u64 for every UI entity. not sure if it's a problem or not

[ickshonpe]

Yeah agree it's not ideal but we have really fat UI nodes already so it shouldn't make much difference. In contrast, each leaked leaked Taffy node is about a thousand bytes. I thought about adding another EntityHashMap instead but the layout module already feels over complicated and I didn't want to even more synchronisation code. This extra NodeId approach is pretty foolproof.