Rebase fixes

crates/bevy_ui/src/layout/debug.rs

@@ -13,17 +13,17 @@ pub fn print_ui_layout_tree(ui_surface: &UiSurface) {
         .entity_to_taffy
         .iter()
         .map(|(&ui_entity, &taffy_node)| (taffy_node, ui_entity))
-        .collect::<HashMap<Node, Entity>>();
+        .collect::<HashMap<NodeId, Entity>>();
     for (&camera_entity, root_node_set) in ui_surface.camera_root_nodes.iter() {
         bevy_utils::tracing::info!("Layout tree for camera entity: {camera_entity}");
         for &root_node_entity in root_node_set.iter() {
             let Some(implicit_viewport_node) = ui_surface
                 .root_node_data
                 .get(&root_node_entity)
                 .map(|rnd| rnd.implicit_viewport_node)
-            else {
-                continue;
-            };
+                else {
+                    continue;
+                };
             let mut out = String::new();
             print_node(
                 ui_surface,
@@ -98,4 +98,4 @@ fn print_node(
             acc,
         );
     }
-}
+}

crates/bevy_ui/src/layout/mod.rs

@@ -193,11 +193,13 @@ pub fn ui_layout_system(
 
     // When a `ContentSize` component is removed from an entity, we need to remove the measure from the corresponding taffy node.
     for entity in removed_components.removed_content_sizes.read() {
-        ui_surface.try_remove_measure(entity);
+        ui_surface.try_remove_node_context(entity);
     }
-    for (entity, mut content_size) in &mut measure_query {
-        if let Some(measure_func) = content_size.measure_func.take() {
-            ui_surface.try_update_measure(entity, measure_func);
+    for (entity, _, mut content_size_option, _) in &mut style_query {
+        if let Some(ref mut content_size) = content_size_option {
+            if let Some(measure_func) = content_size.measure.take() {
+                ui_surface.update_node_context(entity, measure_func);
+            }
         }
     }
     // When a root node is added as a child to another ui node

crates/bevy_ui/src/layout/ui_surface.rs

@@ -49,7 +49,7 @@ pub struct RootNodeData {
     ///
     /// This must be manually removed on `Entity` despawn
     /// or else it will survive in the taffy tree with no references
-    pub(super) implicit_viewport_node: taffy::node::Node,
+    pub(super) implicit_viewport_node: taffy::NodeId,
 }
 
 #[derive(Resource)]
@@ -354,7 +354,7 @@ without UI components as a child of an entity with UI components, results may be
             .unwrap_or_else(|| unreachable!())
     }
 
-    /// Set the ui node entities without a [`Parent`] as children to the root node in the taffy layout.
+    /// Set the ui node entities without a [`bevy_hierarchy::Parent`] as children to the root node in the taffy layout.
     pub fn set_camera_children(
         &mut self,
         camera_entity: Entity,
@@ -423,9 +423,29 @@ without UI components as a child of an entity with UI components, results may be
             }
 
             self.taffy
-                .compute_layout(
+                .compute_layout_with_measure(
                     root_node_data.implicit_viewport_node,
                     available_space,
+                    |known_dimensions: taffy::Size<Option<f32>>,
+                     available_space: taffy::Size<taffy::AvailableSpace>,
+                     _node_id: taffy::NodeId,
+                     context: Option<&mut NodeMeasure>|
+                     -> taffy::Size<f32> {
+                        context
+                            .map(|ctx| {
+                                let size = ctx.measure(
+                                    known_dimensions.width,
+                                    known_dimensions.height,
+                                    available_space.width,
+                                    available_space.height,
+                                );
+                                taffy::Size {
+                                    width: size.x,
+                                    height: size.y,
+                                }
+                            })
+                            .unwrap_or(taffy::Size::ZERO)
+                    },
                 )
                 .unwrap();
         }
@@ -500,14 +520,19 @@ mod tests {
         max_size: 1.0,
     };
 
-    trait IsRootNodeDataValid {
-        fn is_root_node_data_valid(&self, root_node_data: &RootNodeData) -> bool;
+    trait HasValidRootNodeData {
+        fn has_valid_root_node_data(&self, root_node_entity: &Entity) -> bool;
     }
 
-    impl IsRootNodeDataValid for TaffyTree<NodeMeasure> {
-        fn is_root_node_data_valid(&self, root_node_pair: &RootNodeData) -> bool {
-            self.parent(root_node_pair.user_root_node)
-                == Some(root_node_pair.implicit_viewport_node)
+    impl HasValidRootNodeData for UiSurface {
+        fn has_valid_root_node_data(&self, root_node_entity: &Entity) -> bool {
+            let Some(&taffy_node) = self.entity_to_taffy.get(root_node_entity) else {
+                return false;
+            };
+            let Some(root_node_data) = self.root_node_data.get(root_node_entity) else {
+                return false;
+            };
+            self.taffy.parent(taffy_node) == Some(root_node_data.implicit_viewport_node)
         }
     }
 
@@ -538,10 +563,7 @@ mod tests {
         assert_eq!(ui_surface.taffy.total_node_count(), 2);
 
         // root node data should now exist
-        let root_node_data = ui_surface
-            .get_root_node_data(root_node_entity)
-            .expect("expected root node data");
-        assert!(ui_surface.taffy.is_root_node_data_valid(root_node_data));
+        assert!(ui_surface.has_valid_root_node_data(&root_node_entity));
 
         // test duplicate insert 2
         ui_surface.upsert_node(&TEST_LAYOUT_CONTEXT, root_node_entity, &style, None);
@@ -550,10 +572,7 @@ mod tests {
         assert_eq!(ui_surface.taffy.total_node_count(), 2);
 
         // root node data should be unaffected
-        let root_node_data = ui_surface
-            .get_root_node_data(root_node_entity)
-            .expect("expected root node data");
-        assert!(ui_surface.taffy.is_root_node_data_valid(root_node_data));
+        assert!(ui_surface.has_valid_root_node_data(&root_node_entity));
     }
 
     #[test]