MeasureFunc improvements

crates/bevy_ui/src/layout/mod.rs

@@ -1,13 +1,14 @@
 mod convert;
 
-use crate::{CalculatedSize, Node, Style, UiScale};
+use crate::{ContentSize, Node, Style, UiScale};
 use bevy_ecs::{
     change_detection::DetectChanges,
     entity::Entity,
     event::EventReader,
-    query::{Changed, Or, With, Without},
+    query::{Changed, With, Without},
     removal_detection::RemovedComponents,
     system::{Query, Res, ResMut, Resource},
+    world::Ref,
 };
 use bevy_hierarchy::{Children, Parent};
 use bevy_log::warn;
@@ -16,11 +17,7 @@ use bevy_transform::components::Transform;
 use bevy_utils::HashMap;
 use bevy_window::{PrimaryWindow, Window, WindowResolution, WindowScaleFactorChanged};
 use std::fmt;
-use taffy::{
-    prelude::{AvailableSpace, Size},
-    style_helpers::TaffyMaxContent,
-    Taffy,
-};
+use taffy::{prelude::Size, style_helpers::TaffyMaxContent, Taffy};
 
 pub struct LayoutContext {
     pub scale_factor: f64,
@@ -75,6 +72,8 @@ impl Default for UiSurface {
 }
 
 impl UiSurface {
+    /// Retrieves the taffy node corresponding to given entity exists, or inserts a new taffy node into the layout if no corresponding node exists.
+    /// Then convert the given `Style` and use it update the taffy node's style.
     pub fn upsert_node(&mut self, entity: Entity, style: &Style, context: &LayoutContext) {
         let mut added = false;
         let taffy = &mut self.taffy;
@@ -90,43 +89,13 @@ impl UiSurface {
         }
     }
 
-    pub fn upsert_leaf(
-        &mut self,
-        entity: Entity,
-        style: &Style,
-        calculated_size: &CalculatedSize,
-        context: &LayoutContext,
-    ) {
-        let taffy = &mut self.taffy;
-        let taffy_style = convert::from_style(context, style);
-        let measure = calculated_size.measure.dyn_clone();
-        let measure_func = taffy::node::MeasureFunc::Boxed(Box::new(
-            move |constraints: Size<Option<f32>>, available: Size<AvailableSpace>| {
-                let size = measure.measure(
-                    constraints.width,
-                    constraints.height,
-                    available.width,
-                    available.height,
-                );
-                taffy::geometry::Size {
-                    width: size.x,
-                    height: size.y,
-                }
-            },
-        ));
-        if let Some(taffy_node) = self.entity_to_taffy.get(&entity) {
-            self.taffy.set_style(*taffy_node, taffy_style).unwrap();
-            self.taffy
-                .set_measure(*taffy_node, Some(measure_func))
-                .unwrap();
-        } else {
-            let taffy_node = taffy
-                .new_leaf_with_measure(taffy_style, measure_func)
-                .unwrap();
-            self.entity_to_taffy.insert(entity, taffy_node);
-        }
+    /// Update the `MeasureFunc` of the taffy node corresponding to the given [`Entity`].
+    pub fn update_measure(&mut self, entity: Entity, measure_func: taffy::node::MeasureFunc) {
+        let taffy_node = self.entity_to_taffy.get(&entity).unwrap();
+        self.taffy.set_measure(*taffy_node, Some(measure_func)).ok();
     }
 
+    /// Update the children of the taffy node corresponding to the given [`Entity`].
     pub fn update_children(&mut self, entity: Entity, children: &Children) {
         let mut taffy_children = Vec::with_capacity(children.len());
         for child in children {
@@ -160,6 +129,7 @@ without UI components as a child of an entity with UI components, results may be
         }
     }
 
+    /// Retrieve or insert the root layout node and update its size to match the size of the window.
     pub fn update_window(&mut self, window: Entity, window_resolution: &WindowResolution) {
         let taffy = &mut self.taffy;
         let node = self
@@ -185,6 +155,7 @@ without UI components as a child of an entity with UI components, results may be
             .unwrap();
     }
 
+    /// Set the ui node entities without a [`Parent`] as children to the root node in the taffy layout.
     pub fn set_window_children(
         &mut self,
         parent_window: Entity,
@@ -197,6 +168,7 @@ without UI components as a child of an entity with UI components, results may be
         self.taffy.set_children(*taffy_node, &child_nodes).unwrap();
     }
 
+    /// Compute the layout for each window entity's corresponding root node in the layout.
     pub fn compute_window_layouts(&mut self) {
         for window_node in self.window_nodes.values() {
             self.taffy
@@ -214,6 +186,8 @@ without UI components as a child of an entity with UI components, results may be
         }
     }
 
+    /// Get the layout geometry for the taffy node corresponding to the ui node [`Entity`].
+    /// Does not compute the layout geometry, `compute_window_layouts` should be run before using this function.
     pub fn get_layout(&self, entity: Entity) -> Result<&taffy::layout::Layout, LayoutError> {
         if let Some(taffy_node) = self.entity_to_taffy.get(&entity) {
             self.taffy
@@ -235,6 +209,7 @@ pub enum LayoutError {
     TaffyError(taffy::error::TaffyError),
 }
 
+/// Updates the UI's layout tree, computes the new layout geometry and then updates the sizes and transforms of all the UI nodes.
 #[allow(clippy::too_many_arguments)]
 pub fn ui_layout_system(
     primary_window: Query<(Entity, &Window), With<PrimaryWindow>>,
@@ -244,18 +219,11 @@ pub fn ui_layout_system(
     mut resize_events: EventReader<bevy_window::WindowResized>,
     mut ui_surface: ResMut<UiSurface>,
     root_node_query: Query<Entity, (With<Node>, Without<Parent>)>,
-    full_node_query: Query<(Entity, &Style, Option<&CalculatedSize>), With<Node>>,
-    changed_style_query: Query<
-        (Entity, &Style),
-        (With<Node>, Without<CalculatedSize>, Changed<Style>),
-    >,
-    changed_size_query: Query<
-        (Entity, &Style, &CalculatedSize),
-        (With<Node>, Or<(Changed<CalculatedSize>, Changed<Style>)>),
-    >,
+    style_query: Query<(Entity, Ref<Style>), With<Node>>,
+    mut measure_query: Query<(Entity, &mut ContentSize)>,
     children_query: Query<(Entity, &Children), (With<Node>, Changed<Children>)>,
     mut removed_children: RemovedComponents<Children>,
-    mut removed_calculated_sizes: RemovedComponents<CalculatedSize>,
+    mut removed_content_sizes: RemovedComponents<ContentSize>,
     mut node_transform_query: Query<(Entity, &mut Node, &mut Transform, Option<&Parent>)>,
     mut removed_nodes: RemovedComponents<Node>,
 ) {
@@ -285,35 +253,34 @@ pub fn ui_layout_system(
     }
 
     let scale_factor = logical_to_physical_factor * ui_scale.scale;
+
     let layout_context = LayoutContext::new(scale_factor, physical_size);
 
     if !scale_factor_events.is_empty() || ui_scale.is_changed() || resized {
         scale_factor_events.clear();
         // update all nodes
-        for (entity, style, calculated_size) in &full_node_query {
-            if let Some(calculated_size) = calculated_size {
-                ui_surface.upsert_leaf(entity, style, calculated_size, &layout_context);
-            } else {
-                ui_surface.upsert_node(entity, style, &layout_context);
-            }
+        for (entity, style) in style_query.iter() {
+            ui_surface.upsert_node(entity, &style, &layout_context);
         }
     } else {
-        // update changed nodes without a calculated size
-        for (entity, style) in changed_style_query.iter() {
-            ui_surface.upsert_node(entity, style, &layout_context);
+        for (entity, style) in style_query.iter() {
+            if style.is_changed() {
+                ui_surface.upsert_node(entity, &style, &layout_context);
+            }
         }
+    }
 
-        // update changed nodes with a calculated size
-        for (entity, style, calculated_size) in changed_size_query.iter() {
-            ui_surface.upsert_leaf(entity, style, calculated_size, &layout_context);
+    for (entity, mut content_size) in measure_query.iter_mut() {
+        if let Some(measure_func) = content_size.measure_func.take() {
+            ui_surface.update_measure(entity, measure_func);
         }
     }
 
     // clean up removed nodes
     ui_surface.remove_entities(removed_nodes.iter());
 
-    // When a `CalculatedSize` component is removed from an entity, we need to remove the measure from the corresponding taffy node.
-    for entity in removed_calculated_sizes.iter() {
+    // When a `ContentSize` component is removed from an entity, we need to remove the measure from the corresponding taffy node.
+    for entity in removed_content_sizes.iter() {
         ui_surface.try_remove_measure(entity);
     }
 

crates/bevy_ui/src/lib.rs

@@ -87,7 +87,7 @@ impl Plugin for UiPlugin {
             .register_type::<AlignContent>()
             .register_type::<AlignItems>()
             .register_type::<AlignSelf>()
-            .register_type::<CalculatedSize>()
+            .register_type::<ContentSize>()
             .register_type::<Direction>()
             .register_type::<Display>()
             .register_type::<FlexDirection>()
@@ -142,7 +142,7 @@ impl Plugin for UiPlugin {
         #[cfg(feature = "bevy_text")]
         app.add_plugin(accessibility::AccessibilityPlugin);
         app.add_systems(PostUpdate, {
-            let system = widget::update_image_calculated_size_system.before(UiSystem::Layout);
+            let system = widget::update_image_content_size_system.before(UiSystem::Layout);
             // Potential conflicts: `Assets<Image>`
             // They run independently since `widget::image_node_system` will only ever observe
             // its own UiImage, and `widget::text_system` & `bevy_text::update_text2d_layout`

crates/bevy_ui/src/measurement.rs

@@ -1,12 +1,13 @@
 use bevy_ecs::prelude::Component;
+use bevy_ecs::reflect::ReflectComponent;
 use bevy_math::Vec2;
 use bevy_reflect::Reflect;
 use std::fmt::Formatter;
 pub use taffy::style::AvailableSpace;
 
-impl std::fmt::Debug for CalculatedSize {
+impl std::fmt::Debug for ContentSize {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("CalculatedSize").finish()
+        f.debug_struct("ContentSize").finish()
     }
 }
 
@@ -21,9 +22,6 @@ pub trait Measure: Send + Sync + 'static {
         available_width: AvailableSpace,
         available_height: AvailableSpace,
     ) -> Vec2;
-
-    /// Clone and box self.
-    fn dyn_clone(&self) -> Box<dyn Measure>;
 }
 
 /// A `FixedMeasure` is a `Measure` that ignores all constraints and
@@ -43,35 +41,42 @@ impl Measure for FixedMeasure {
     ) -> Vec2 {
         self.size
     }
-
-    fn dyn_clone(&self) -> Box<dyn Measure> {
-        Box::new(self.clone())
-    }
 }
 
-/// A node with a `CalculatedSize` component is a node where its size
+/// A node with a `ContentSize` component is a node where its size
 /// is based on its content.
 #[derive(Component, Reflect)]
-pub struct CalculatedSize {
+#[reflect(Component)]
+pub struct ContentSize {
     /// The `Measure` used to compute the intrinsic size
     #[reflect(ignore)]
-    pub measure: Box<dyn Measure>,
+    pub(crate) measure_func: Option<taffy::node::MeasureFunc>,
 }
 
-#[allow(clippy::derivable_impls)]
-impl Default for CalculatedSize {
-    fn default() -> Self {
-        Self {
-            // Default `FixedMeasure` always returns zero size.
-            measure: Box::<FixedMeasure>::default(),
-        }
+impl ContentSize {
+    /// Set a `Measure` for this function
+    pub fn set(&mut self, measure: impl Measure) {
+        let measure_func =
+            move |size: taffy::prelude::Size<Option<f32>>,
+                  available: taffy::prelude::Size<AvailableSpace>| {
+                let size =
+                    measure.measure(size.width, size.height, available.width, available.height);
+                taffy::prelude::Size {
+                    width: size.x,
+                    height: size.y,
+                }
+            };
+        self.measure_func = Some(taffy::node::MeasureFunc::Boxed(Box::new(measure_func)));
     }
 }
 
-impl Clone for CalculatedSize {
-    fn clone(&self) -> Self {
+#[allow(clippy::derivable_impls)]
+impl Default for ContentSize {
+    fn default() -> Self {
         Self {
-            measure: self.measure.dyn_clone(),
+            measure_func: Some(taffy::node::MeasureFunc::Raw(|_, _| {
+                taffy::prelude::Size::ZERO
+            })),
         }
     }
 }

crates/bevy_ui/src/node_bundles.rs

@@ -2,7 +2,7 @@
 
 use crate::{
     widget::{Button, UiImageSize},
-    BackgroundColor, CalculatedSize, FocusPolicy, Interaction, Node, Style, UiImage, ZIndex,
+    BackgroundColor, ContentSize, FocusPolicy, Interaction, Node, Style, UiImage, ZIndex,
 };
 use bevy_ecs::bundle::Bundle;
 use bevy_render::{
@@ -63,7 +63,7 @@ impl Default for NodeBundle {
 }
 
 /// A UI node that is an image
-#[derive(Bundle, Clone, Debug, Default)]
+#[derive(Bundle, Debug, Default)]
 pub struct ImageBundle {
     /// Describes the logical size of the node
     ///
@@ -74,7 +74,7 @@ pub struct ImageBundle {
     /// In some cases these styles also affect how the node drawn/painted.
     pub style: Style,
     /// The calculated size based on the given image
-    pub calculated_size: CalculatedSize,
+    pub calculated_size: ContentSize,
     /// The background color, which serves as a "fill" for this node
     ///
     /// Combines with `UiImage` to tint the provided image.
@@ -107,7 +107,7 @@ pub struct ImageBundle {
 
 #[cfg(feature = "bevy_text")]
 /// A UI node that is text
-#[derive(Bundle, Clone, Debug)]
+#[derive(Bundle, Debug)]
 pub struct TextBundle {
     /// Describes the logical size of the node
     pub node: Node,
@@ -119,7 +119,7 @@ pub struct TextBundle {
     /// Text layout information
     pub text_layout_info: TextLayoutInfo,
     /// The calculated size based on the given image
-    pub calculated_size: CalculatedSize,
+    pub calculated_size: ContentSize,
     /// Whether this node should block interaction with lower nodes
     pub focus_policy: FocusPolicy,
     /// The transform of the node

crates/bevy_ui/src/widget/image.rs

@@ -1,4 +1,4 @@
-use crate::{measurement::AvailableSpace, CalculatedSize, Measure, Node, UiImage};
+use crate::{measurement::AvailableSpace, ContentSize, Measure, Node, UiImage};
 use bevy_asset::Assets;
 #[cfg(feature = "bevy_text")]
 use bevy_ecs::query::Without;
@@ -58,25 +58,21 @@ impl Measure for ImageMeasure {
         }
         size
     }
-
-    fn dyn_clone(&self) -> Box<dyn Measure> {
-        Box::new(self.clone())
-    }
 }
 
-/// Updates calculated size of the node based on the image provided
-pub fn update_image_calculated_size_system(
+/// Updates content size of the node based on the image provided
+pub fn update_image_content_size_system(
     textures: Res<Assets<Image>>,
     #[cfg(feature = "bevy_text")] mut query: Query<
-        (&mut CalculatedSize, &UiImage, &mut UiImageSize),
+        (&mut ContentSize, &UiImage, &mut UiImageSize),
         (With<Node>, Without<Text>),
     >,
     #[cfg(not(feature = "bevy_text"))] mut query: Query<
-        (&mut CalculatedSize, &UiImage, &mut UiImageSize),
+        (&mut ContentSize, &UiImage, &mut UiImageSize),
         With<Node>,
     >,
 ) {
-    for (mut calculated_size, image, mut image_size) in &mut query {
+    for (mut content_size, image, mut image_size) in &mut query {
         if let Some(texture) = textures.get(&image.texture) {
             let size = Vec2::new(
                 texture.texture_descriptor.size.width as f32,
@@ -85,7 +81,7 @@ pub fn update_image_calculated_size_system(
             // Update only if size has changed to avoid needless layout calculations
             if size != image_size.size {
                 image_size.size = size;
-                calculated_size.measure = Box::new(ImageMeasure { size });
+                content_size.set(ImageMeasure { size });
             }
         }
     }