From 96b9d0a7e203bc55aff4f8b14ac8a51a6945fea5 Mon Sep 17 00:00:00 2001
From: Nico Burns <nico@nicoburns.com>
Date: Wed, 1 May 2024 02:13:17 +1200
Subject: [PATCH] Upgrade to Taffy 0.4 (#10690)

# Objective

- Enables support for `Display::Block`
- Enables support for `Overflow::Hidden`
- Allows for cleaner integration with text, image and other content
layout.
- Unblocks https://github.com/bevyengine/bevy/pull/8104
- Unlocks the possibility of Bevy creating a custom layout tree over
which Taffy operates.
- Enables #8808 / #10193 to remove a Mutex around the font system.

## Todo

- [x] ~Fix rendering of text/images to account for padding/border on
nodes (should size/position to content box rather than border box)~ In
order get this into a mergeable state this PR instead zeroes out
padding/border when syncing leaf node styles into Taffy to preserve the
existing behaviour. https://github.com/bevyengine/bevy/issues/6879 can
be fixed in a followup PR.

## Solution

- Update the version of Taffy
- Update code to work with the new version

Note: Taffy 0.4 has not yet been released. This PR is being created in
advance of the release to ensure that there are no blockers to upgrading
once the release occurs.

---

## Changelog

- Bevy now supports the `Display::Block` and `Overflow::Hidden` styles.
---
 crates/bevy_ui/Cargo.toml               |   2 +-
 crates/bevy_ui/src/layout/convert.rs    | 108 +++++++++++++++---------
 crates/bevy_ui/src/layout/debug.rs      |  14 +--
 crates/bevy_ui/src/layout/mod.rs        |  60 +++++++------
 crates/bevy_ui/src/layout/ui_surface.rs | 103 ++++++++++++++++------
 crates/bevy_ui/src/measurement.rs       |  59 ++++++++++---
 crates/bevy_ui/src/ui_node.rs           |  10 ++-
 crates/bevy_ui/src/widget/image.rs      |   9 +-
 crates/bevy_ui/src/widget/text.rs       |   6 +-
 examples/ui/display_and_visibility.rs   |   2 +-
 10 files changed, 250 insertions(+), 123 deletions(-)

diff --git a/crates/bevy_ui/Cargo.toml b/crates/bevy_ui/Cargo.toml
index d7882c7da4d38..4e47819a6da6c 100644
--- a/crates/bevy_ui/Cargo.toml
+++ b/crates/bevy_ui/Cargo.toml
@@ -31,7 +31,7 @@ bevy_window = { path = "../bevy_window", version = "0.14.0-dev" }
 bevy_utils = { path = "../bevy_utils", version = "0.14.0-dev" }
 
 # other
-taffy = { version = "0.3.10" }
+taffy = { version = "0.4" }
 serde = { version = "1", features = ["derive"], optional = true }
 bytemuck = { version = "1.5", features = ["derive"] }
 thiserror = "1.0.0"
diff --git a/crates/bevy_ui/src/layout/convert.rs b/crates/bevy_ui/src/layout/convert.rs
index ce93ae806a240..2cc7c3b32c0f3 100644
--- a/crates/bevy_ui/src/layout/convert.rs
+++ b/crates/bevy_ui/src/layout/convert.rs
@@ -3,8 +3,8 @@ use taffy::style_helpers;
 use crate::{
     AlignContent, AlignItems, AlignSelf, Display, FlexDirection, FlexWrap, GridAutoFlow,
     GridPlacement, GridTrack, GridTrackRepetition, JustifyContent, JustifyItems, JustifySelf,
-    MaxTrackSizingFunction, MinTrackSizingFunction, PositionType, RepeatedGridTrack, Style, UiRect,
-    Val,
+    MaxTrackSizingFunction, MinTrackSizingFunction, OverflowAxis, PositionType, RepeatedGridTrack,
+    Style, UiRect, Val,
 };
 
 use super::LayoutContext;
@@ -18,31 +18,31 @@ impl Val {
             Val::Auto => taffy::style::LengthPercentageAuto::Auto,
             Val::Percent(value) => taffy::style::LengthPercentageAuto::Percent(value / 100.),
             Val::Px(value) => {
-                taffy::style::LengthPercentageAuto::Points(context.scale_factor * value)
+                taffy::style::LengthPercentageAuto::Length(context.scale_factor * value)
             }
             Val::VMin(value) => {
-                taffy::style::LengthPercentageAuto::Points(context.min_size * value / 100.)
+                taffy::style::LengthPercentageAuto::Length(context.min_size * value / 100.)
             }
             Val::VMax(value) => {
-                taffy::style::LengthPercentageAuto::Points(context.max_size * value / 100.)
+                taffy::style::LengthPercentageAuto::Length(context.max_size * value / 100.)
             }
             Val::Vw(value) => {
-                taffy::style::LengthPercentageAuto::Points(context.physical_size.x * value / 100.)
+                taffy::style::LengthPercentageAuto::Length(context.physical_size.x * value / 100.)
             }
             Val::Vh(value) => {
-                taffy::style::LengthPercentageAuto::Points(context.physical_size.y * value / 100.)
+                taffy::style::LengthPercentageAuto::Length(context.physical_size.y * value / 100.)
             }
         }
     }
 
     fn into_length_percentage(self, context: &LayoutContext) -> taffy::style::LengthPercentage {
         match self.into_length_percentage_auto(context) {
-            taffy::style::LengthPercentageAuto::Auto => taffy::style::LengthPercentage::Points(0.0),
+            taffy::style::LengthPercentageAuto::Auto => taffy::style::LengthPercentage::Length(0.0),
             taffy::style::LengthPercentageAuto::Percent(value) => {
                 taffy::style::LengthPercentage::Percent(value)
             }
-            taffy::style::LengthPercentageAuto::Points(value) => {
-                taffy::style::LengthPercentage::Points(value)
+            taffy::style::LengthPercentageAuto::Length(value) => {
+                taffy::style::LengthPercentage::Length(value)
             }
         }
     }
@@ -63,9 +63,18 @@ impl UiRect {
     }
 }
 
-pub fn from_style(context: &LayoutContext, style: &Style) -> taffy::style::Style {
+pub fn from_style(
+    context: &LayoutContext,
+    style: &Style,
+    ignore_padding_and_border: bool,
+) -> taffy::style::Style {
     taffy::style::Style {
         display: style.display.into(),
+        overflow: taffy::Point {
+            x: style.overflow.x.into(),
+            y: style.overflow.y.into(),
+        },
+        scrollbar_width: 0.0,
         position: style.position_type.into(),
         flex_direction: style.flex_direction.into(),
         flex_wrap: style.flex_wrap.into(),
@@ -75,7 +84,7 @@ pub fn from_style(context: &LayoutContext, style: &Style) -> taffy::style::Style
         justify_self: style.justify_self.into(),
         align_content: style.align_content.into(),
         justify_content: style.justify_content.into(),
-        inset: taffy::prelude::Rect {
+        inset: taffy::Rect {
             left: style.left.into_length_percentage_auto(context),
             right: style.right.into_length_percentage_auto(context),
             top: style.top.into_length_percentage_auto(context),
@@ -84,29 +93,41 @@ pub fn from_style(context: &LayoutContext, style: &Style) -> taffy::style::Style
         margin: style
             .margin
             .map_to_taffy_rect(|m| m.into_length_percentage_auto(context)),
-        padding: style
-            .padding
-            .map_to_taffy_rect(|m| m.into_length_percentage(context)),
-        border: style
-            .border
-            .map_to_taffy_rect(|m| m.into_length_percentage(context)),
+        // Ignore padding for leaf nodes as it isn't implemented in the rendering engine.
+        // TODO: Implement rendering of padding for leaf nodes
+        padding: if ignore_padding_and_border {
+            taffy::Rect::zero()
+        } else {
+            style
+                .padding
+                .map_to_taffy_rect(|m| m.into_length_percentage(context))
+        },
+        // Ignore border for leaf nodes as it isn't implemented in the rendering engine.
+        // TODO: Implement rendering of border for leaf nodes
+        border: if ignore_padding_and_border {
+            taffy::Rect::zero()
+        } else {
+            style
+                .border
+                .map_to_taffy_rect(|m| m.into_length_percentage(context))
+        },
         flex_grow: style.flex_grow,
         flex_shrink: style.flex_shrink,
         flex_basis: style.flex_basis.into_dimension(context),
-        size: taffy::prelude::Size {
+        size: taffy::Size {
             width: style.width.into_dimension(context),
             height: style.height.into_dimension(context),
         },
-        min_size: taffy::prelude::Size {
+        min_size: taffy::Size {
             width: style.min_width.into_dimension(context),
             height: style.min_height.into_dimension(context),
         },
-        max_size: taffy::prelude::Size {
+        max_size: taffy::Size {
             width: style.max_width.into_dimension(context),
             height: style.max_height.into_dimension(context),
         },
         aspect_ratio: style.aspect_ratio,
-        gap: taffy::prelude::Size {
+        gap: taffy::Size {
             width: style.column_gap.into_length_percentage(context),
             height: style.row_gap.into_length_percentage(context),
         },
@@ -231,11 +252,22 @@ impl From<Display> for taffy::style::Display {
         match value {
             Display::Flex => taffy::style::Display::Flex,
             Display::Grid => taffy::style::Display::Grid,
+            Display::Block => taffy::style::Display::Block,
             Display::None => taffy::style::Display::None,
         }
     }
 }
 
+impl From<OverflowAxis> for taffy::style::Overflow {
+    fn from(value: OverflowAxis) -> Self {
+        match value {
+            OverflowAxis::Visible => taffy::style::Overflow::Visible,
+            OverflowAxis::Clip => taffy::style::Overflow::Clip,
+            OverflowAxis::Hidden => taffy::style::Overflow::Hidden,
+        }
+    }
+}
+
 impl From<FlexDirection> for taffy::style::FlexDirection {
     fn from(value: FlexDirection) -> Self {
         match value {
@@ -489,7 +521,7 @@ mod tests {
             grid_row: GridPlacement::span(3),
         };
         let viewport_values = LayoutContext::new(1.0, bevy_math::Vec2::new(800., 600.));
-        let taffy_style = from_style(&viewport_values, &bevy_style);
+        let taffy_style = from_style(&viewport_values, &bevy_style, false);
         assert_eq!(taffy_style.display, taffy::style::Display::Flex);
         assert_eq!(taffy_style.position, taffy::style::Position::Absolute);
         assert_eq!(
@@ -502,7 +534,7 @@ mod tests {
         );
         assert_eq!(
             taffy_style.inset.top,
-            taffy::style::LengthPercentageAuto::Points(12.)
+            taffy::style::LengthPercentageAuto::Length(12.)
         );
         assert_eq!(
             taffy_style.inset.bottom,
@@ -537,7 +569,7 @@ mod tests {
         );
         assert_eq!(
             taffy_style.margin.right,
-            taffy::style::LengthPercentageAuto::Points(10.)
+            taffy::style::LengthPercentageAuto::Length(10.)
         );
         assert_eq!(
             taffy_style.margin.top,
@@ -553,7 +585,7 @@ mod tests {
         );
         assert_eq!(
             taffy_style.padding.right,
-            taffy::style::LengthPercentage::Points(21.)
+            taffy::style::LengthPercentage::Length(21.)
         );
         assert_eq!(
             taffy_style.padding.top,
@@ -565,7 +597,7 @@ mod tests {
         );
         assert_eq!(
             taffy_style.border.left,
-            taffy::style::LengthPercentage::Points(14.)
+            taffy::style::LengthPercentage::Length(14.)
         );
         assert_eq!(
             taffy_style.border.right,
@@ -594,18 +626,18 @@ mod tests {
         );
         assert_eq!(
             taffy_style.grid_template_rows,
-            vec![sh::points(10.0), sh::percent(0.5), sh::fr(1.0)]
+            vec![sh::length(10.0), sh::percent(0.5), sh::fr(1.0)]
         );
         assert_eq!(
             taffy_style.grid_template_columns,
-            vec![sh::repeat(5, vec![sh::points(10.0)])]
+            vec![sh::repeat(5, vec![sh::length(10.0)])]
         );
         assert_eq!(
             taffy_style.grid_auto_rows,
             vec![
-                sh::fit_content(taffy::style::LengthPercentage::Points(10.0)),
+                sh::fit_content(taffy::style::LengthPercentage::Length(10.0)),
                 sh::fit_content(taffy::style::LengthPercentage::Percent(0.25)),
-                sh::minmax(sh::points(0.0), sh::fr(2.0)),
+                sh::minmax(sh::length(0.0), sh::fr(2.0)),
             ]
         );
         assert_eq!(
@@ -627,17 +659,17 @@ mod tests {
         use taffy::style::LengthPercentage;
         let context = LayoutContext::new(2.0, bevy_math::Vec2::new(800., 600.));
         let cases = [
-            (Val::Auto, LengthPercentage::Points(0.)),
+            (Val::Auto, LengthPercentage::Length(0.)),
             (Val::Percent(1.), LengthPercentage::Percent(0.01)),
-            (Val::Px(1.), LengthPercentage::Points(2.)),
-            (Val::Vw(1.), LengthPercentage::Points(8.)),
-            (Val::Vh(1.), LengthPercentage::Points(6.)),
-            (Val::VMin(2.), LengthPercentage::Points(12.)),
-            (Val::VMax(2.), LengthPercentage::Points(16.)),
+            (Val::Px(1.), LengthPercentage::Length(2.)),
+            (Val::Vw(1.), LengthPercentage::Length(8.)),
+            (Val::Vh(1.), LengthPercentage::Length(6.)),
+            (Val::VMin(2.), LengthPercentage::Length(12.)),
+            (Val::VMax(2.), LengthPercentage::Length(16.)),
         ];
         for (val, length) in cases {
             assert!(match (val.into_length_percentage(&context), length) {
-                (LengthPercentage::Points(a), LengthPercentage::Points(b))
+                (LengthPercentage::Length(a), LengthPercentage::Length(b))
                 | (LengthPercentage::Percent(a), LengthPercentage::Percent(b)) =>
                     (a - b).abs() < 0.0001,
                 _ => false,
diff --git a/crates/bevy_ui/src/layout/debug.rs b/crates/bevy_ui/src/layout/debug.rs
index 67e7a205b2915..47b02396a6816 100644
--- a/crates/bevy_ui/src/layout/debug.rs
+++ b/crates/bevy_ui/src/layout/debug.rs
@@ -1,7 +1,6 @@
 use std::fmt::Write;
 
-use taffy::prelude::Node;
-use taffy::tree::LayoutTree;
+use taffy::{NodeId, TraversePartialTree};
 
 use bevy_ecs::prelude::Entity;
 use bevy_utils::HashMap;
@@ -10,7 +9,7 @@ use crate::layout::ui_surface::UiSurface;
 
 /// Prints a debug representation of the computed layout of the UI layout tree for each window.
 pub fn print_ui_layout_tree(ui_surface: &UiSurface) {
-    let taffy_to_entity: HashMap<Node, Entity> = ui_surface
+    let taffy_to_entity: HashMap<NodeId, Entity> = ui_surface
         .entity_to_taffy
         .iter()
         .map(|(entity, node)| (*node, *entity))
@@ -35,9 +34,9 @@ pub fn print_ui_layout_tree(ui_surface: &UiSurface) {
 /// Recursively navigates the layout tree printing each node's information.
 fn print_node(
     ui_surface: &UiSurface,
-    taffy_to_entity: &HashMap<Node, Entity>,
+    taffy_to_entity: &HashMap<NodeId, Entity>,
     entity: Entity,
-    node: Node,
+    node: NodeId,
     has_sibling: bool,
     lines_string: String,
     acc: &mut String,
@@ -46,13 +45,14 @@ fn print_node(
     let layout = tree.layout(node).unwrap();
     let style = tree.style(node).unwrap();
 
-    let num_children = tree.child_count(node).unwrap();
+    let num_children = tree.child_count(node);
 
     let display_variant = match (num_children, style.display) {
         (_, taffy::style::Display::None) => "NONE",
         (0, _) => "LEAF",
         (_, taffy::style::Display::Flex) => "FLEX",
         (_, taffy::style::Display::Grid) => "GRID",
+        (_, taffy::style::Display::Block) => "BLOCK",
     };
 
     let fork_string = if has_sibling {
@@ -70,7 +70,7 @@ fn print_node(
         y = layout.location.y,
         width = layout.size.width,
         height = layout.size.height,
-        measured = if tree.needs_measure(node) { "measured" } else { "" }
+        measured = if tree.get_node_context(node).is_some() { "measured" } else { "" }
     ).ok();
     let bar = if has_sibling { "│   " } else { "    " };
     let new_string = lines_string + bar;
diff --git a/crates/bevy_ui/src/layout/mod.rs b/crates/bevy_ui/src/layout/mod.rs
index ff82518677ef4..43d767acdecbc 100644
--- a/crates/bevy_ui/src/layout/mod.rs
+++ b/crates/bevy_ui/src/layout/mod.rs
@@ -1,5 +1,6 @@
 use thiserror::Error;
 
+use crate::{ContentSize, DefaultUiCamera, Node, Outline, Style, TargetCamera, UiScale};
 use bevy_ecs::{
     change_detection::{DetectChanges, DetectChangesMut},
     entity::Entity,
@@ -18,8 +19,6 @@ use bevy_utils::{HashMap, HashSet};
 use bevy_window::{PrimaryWindow, Window, WindowScaleFactorChanged};
 use ui_surface::UiSurface;
 
-use crate::{ContentSize, DefaultUiCamera, Node, Outline, Style, TargetCamera, UiScale};
-
 mod convert;
 pub mod debug;
 pub(crate) mod ui_surface;
@@ -60,7 +59,7 @@ pub enum LayoutError {
     #[error("Invalid hierarchy")]
     InvalidHierarchy,
     #[error("Taffy error: {0}")]
-    TaffyError(#[from] taffy::error::TaffyError),
+    TaffyError(#[from] taffy::TaffyError),
 }
 
 #[derive(SystemParam)]
@@ -82,8 +81,15 @@ pub fn ui_layout_system(
     mut resize_events: EventReader<bevy_window::WindowResized>,
     mut ui_surface: ResMut<UiSurface>,
     root_node_query: Query<(Entity, Option<&TargetCamera>), (With<Node>, Without<Parent>)>,
-    style_query: Query<(Entity, Ref<Style>, Option<&TargetCamera>), With<Node>>,
-    mut measure_query: Query<(Entity, &mut ContentSize)>,
+    mut style_query: Query<
+        (
+            Entity,
+            Ref<Style>,
+            Option<&mut ContentSize>,
+            Option<&TargetCamera>,
+        ),
+        With<Node>,
+    >,
     children_query: Query<(Entity, Ref<Children>), With<Node>>,
     just_children_query: Query<&Children>,
     mut removed_components: UiLayoutSystemRemovedComponentParam,
@@ -152,8 +158,13 @@ pub fn ui_layout_system(
         }
     }
 
-    // Resize all nodes
-    for (entity, style, target_camera) in style_query.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_components.removed_content_sizes.read() {
+        ui_surface.try_remove_node_context(entity);
+    }
+
+    // Sync Style and ContentSize to Taffy for all nodes
+    for (entity, style, content_size, target_camera) in style_query.iter_mut() {
         if let Some(camera) =
             camera_with_default(target_camera).and_then(|c| camera_layout_info.get(&c))
         {
@@ -161,29 +172,24 @@ pub fn ui_layout_system(
                 || !scale_factor_events.is_empty()
                 || ui_scale.is_changed()
                 || style.is_changed()
+                || content_size
+                    .as_ref()
+                    .map(|c| c.measure.is_some())
+                    .unwrap_or(false)
             {
                 let layout_context = LayoutContext::new(
                     camera.scale_factor,
                     [camera.size.x as f32, camera.size.y as f32].into(),
                 );
-                ui_surface.upsert_node(entity, &style, &layout_context);
+                let measure = content_size.and_then(|mut c| c.measure.take());
+                ui_surface.upsert_node(&layout_context, entity, &style, measure);
             }
         } else {
-            ui_surface.upsert_node(entity, &Style::default(), &LayoutContext::default());
+            ui_surface.upsert_node(&LayoutContext::DEFAULT, entity, &Style::default(), None);
         }
     }
     scale_factor_events.clear();
 
-    // 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);
-    }
-    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);
-        }
-    }
-
     // clean up removed nodes
     ui_surface.remove_entities(removed_components.removed_nodes.read());
 
@@ -337,7 +343,7 @@ fn round_layout_coords(value: Vec2) -> Vec2 {
 
 #[cfg(test)]
 mod tests {
-    use taffy::tree::LayoutTree;
+    use taffy::TraversePartialTree;
 
     use bevy_asset::AssetEvent;
     use bevy_asset::Assets;
@@ -578,7 +584,7 @@ mod tests {
         let ui_parent_node = ui_surface.entity_to_taffy[&ui_parent_entity];
 
         // `ui_parent_node` shouldn't have any children yet
-        assert_eq!(ui_surface.taffy.child_count(ui_parent_node).unwrap(), 0);
+        assert_eq!(ui_surface.taffy.child_count(ui_parent_node), 0);
 
         let mut ui_child_entities = (0..10)
             .map(|_| {
@@ -597,7 +603,7 @@ mod tests {
             1 + ui_child_entities.len()
         );
         assert_eq!(
-            ui_surface.taffy.child_count(ui_parent_node).unwrap(),
+            ui_surface.taffy.child_count(ui_parent_node),
             ui_child_entities.len()
         );
 
@@ -628,7 +634,7 @@ mod tests {
             1 + ui_child_entities.len()
         );
         assert_eq!(
-            ui_surface.taffy.child_count(ui_parent_node).unwrap(),
+            ui_surface.taffy.child_count(ui_parent_node),
             ui_child_entities.len()
         );
 
@@ -941,8 +947,8 @@ mod tests {
         let ui_surface = world.resource::<UiSurface>();
         let ui_node = ui_surface.entity_to_taffy[&ui_entity];
 
-        // a node with a content size needs to be measured
-        assert!(ui_surface.taffy.needs_measure(ui_node));
+        // a node with a content size should have taffy context
+        assert!(ui_surface.taffy.get_node_context(ui_node).is_some());
         let layout = ui_surface.get_layout(ui_entity).unwrap();
         assert_eq!(layout.size.width, content_size.x);
         assert_eq!(layout.size.height, content_size.y);
@@ -952,8 +958,8 @@ mod tests {
         ui_schedule.run(&mut world);
 
         let ui_surface = world.resource::<UiSurface>();
-        // a node without a content size does not need to be measured
-        assert!(!ui_surface.taffy.needs_measure(ui_node));
+        // a node without a content size should not have taffy context
+        assert!(ui_surface.taffy.get_node_context(ui_node).is_none());
 
         // Without a content size, the node has no width or height constraints so the length of both dimensions is 0.
         let layout = ui_surface.get_layout(ui_entity).unwrap();
diff --git a/crates/bevy_ui/src/layout/ui_surface.rs b/crates/bevy_ui/src/layout/ui_surface.rs
index 74586e8786ac0..2ba6eed374383 100644
--- a/crates/bevy_ui/src/layout/ui_surface.rs
+++ b/crates/bevy_ui/src/layout/ui_surface.rs
@@ -1,7 +1,6 @@
 use std::fmt;
 
-use taffy::prelude::LayoutTree;
-use taffy::Taffy;
+use taffy::TaffyTree;
 
 use bevy_ecs::entity::{Entity, EntityHashMap};
 use bevy_ecs::prelude::Resource;
@@ -11,28 +10,28 @@ use bevy_utils::default;
 use bevy_utils::tracing::warn;
 
 use crate::layout::convert;
-use crate::{LayoutContext, LayoutError, Style};
+use crate::{LayoutContext, LayoutError, Measure, NodeMeasure, Style};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct RootNodePair {
     // The implicit "viewport" node created by Bevy
-    pub(super) implicit_viewport_node: taffy::node::Node,
+    pub(super) implicit_viewport_node: taffy::NodeId,
     // The root (parentless) node specified by the user
-    pub(super) user_root_node: taffy::node::Node,
+    pub(super) user_root_node: taffy::NodeId,
 }
 
 #[derive(Resource)]
 pub struct UiSurface {
-    pub(super) entity_to_taffy: EntityHashMap<taffy::node::Node>,
-    pub(super) camera_entity_to_taffy: EntityHashMap<EntityHashMap<taffy::node::Node>>,
+    pub(super) entity_to_taffy: EntityHashMap<taffy::NodeId>,
+    pub(super) camera_entity_to_taffy: EntityHashMap<EntityHashMap<taffy::NodeId>>,
     pub(super) camera_roots: EntityHashMap<Vec<RootNodePair>>,
-    pub(super) taffy: Taffy,
+    pub(super) taffy: TaffyTree<NodeMeasure>,
 }
 
 fn _assert_send_sync_ui_surface_impl_safe() {
     fn _assert_send_sync<T: Send + Sync>() {}
-    _assert_send_sync::<EntityHashMap<taffy::node::Node>>();
-    _assert_send_sync::<Taffy>();
+    _assert_send_sync::<EntityHashMap<taffy::NodeId>>();
+    _assert_send_sync::<TaffyTree<NodeMeasure>>();
     _assert_send_sync::<UiSurface>();
 }
 
@@ -47,7 +46,7 @@ impl fmt::Debug for UiSurface {
 
 impl Default for UiSurface {
     fn default() -> Self {
-        let mut taffy = Taffy::new();
+        let mut taffy: TaffyTree<NodeMeasure> = TaffyTree::new();
         taffy.disable_rounding();
         Self {
             entity_to_taffy: Default::default(),
@@ -61,30 +60,55 @@ impl Default for UiSurface {
 impl UiSurface {
     /// Retrieves the Taffy node associated with the given UI node entity and updates its style.
     /// If no associated Taffy node exists a new Taffy node is inserted into the Taffy layout.
-    pub fn upsert_node(&mut self, entity: Entity, style: &Style, context: &LayoutContext) {
-        let mut added = false;
+    pub fn upsert_node(
+        &mut self,
+        layout_context: &LayoutContext,
+        entity: Entity,
+        style: &Style,
+        mut new_node_context: Option<NodeMeasure>,
+    ) {
         let taffy = &mut self.taffy;
-        let taffy_node = self.entity_to_taffy.entry(entity).or_insert_with(|| {
+
+        let mut added = false;
+        let taffy_node_id = *self.entity_to_taffy.entry(entity).or_insert_with(|| {
             added = true;
-            taffy.new_leaf(convert::from_style(context, style)).unwrap()
+            if let Some(measure) = new_node_context.take() {
+                taffy
+                    .new_leaf_with_context(
+                        convert::from_style(layout_context, style, true),
+                        measure,
+                    )
+                    .unwrap()
+            } else {
+                taffy
+                    .new_leaf(convert::from_style(layout_context, style, false))
+                    .unwrap()
+            }
         });
 
         if !added {
-            self.taffy
-                .set_style(*taffy_node, convert::from_style(context, style))
+            let has_measure = if new_node_context.is_some() {
+                taffy
+                    .set_node_context(taffy_node_id, new_node_context)
+                    .unwrap();
+                true
+            } else {
+                taffy.get_node_context(taffy_node_id).is_some()
+            };
+
+            taffy
+                .set_style(
+                    taffy_node_id,
+                    convert::from_style(layout_context, style, has_measure),
+                )
                 .unwrap();
         }
     }
 
     /// Update the `MeasureFunc` of the taffy node corresponding to the given [`Entity`] if the node exists.
-    pub fn try_update_measure(
-        &mut self,
-        entity: Entity,
-        measure_func: taffy::node::MeasureFunc,
-    ) -> Option<()> {
+    pub fn update_node_context(&mut self, entity: Entity, context: NodeMeasure) -> Option<()> {
         let taffy_node = self.entity_to_taffy.get(&entity)?;
-
-        self.taffy.set_measure(*taffy_node, Some(measure_func)).ok()
+        self.taffy.set_node_context(*taffy_node, Some(context)).ok()
     }
 
     /// Update the children of the taffy node corresponding to the given [`Entity`].
@@ -115,9 +139,9 @@ without UI components as a child of an entity with UI components, results may be
     }
 
     /// Removes the measure from the entity's taffy node if it exists. Does nothing otherwise.
-    pub fn try_remove_measure(&mut self, entity: Entity) {
+    pub fn try_remove_node_context(&mut self, entity: Entity) {
         if let Some(taffy_node) = self.entity_to_taffy.get(&entity) {
-            self.taffy.set_measure(*taffy_node, None).unwrap();
+            self.taffy.set_node_context(*taffy_node, None).unwrap();
         }
     }
 
@@ -183,7 +207,30 @@ without UI components as a child of an entity with UI components, results may be
         };
         for root_nodes in camera_root_nodes {
             self.taffy
-                .compute_layout(root_nodes.implicit_viewport_node, available_space)
+                .compute_layout_with_measure(
+                    root_nodes.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();
         }
     }
@@ -210,7 +257,7 @@ 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> {
+    pub fn get_layout(&self, entity: Entity) -> Result<&taffy::Layout, LayoutError> {
         if let Some(taffy_node) = self.entity_to_taffy.get(&entity) {
             self.taffy
                 .layout(*taffy_node)
diff --git a/crates/bevy_ui/src/measurement.rs b/crates/bevy_ui/src/measurement.rs
index 1457dd6a3af8d..b482bd01aa759 100644
--- a/crates/bevy_ui/src/measurement.rs
+++ b/crates/bevy_ui/src/measurement.rs
@@ -4,7 +4,11 @@ use bevy_math::Vec2;
 use bevy_reflect::{std_traits::ReflectDefault, Reflect};
 use std::fmt::Formatter;
 pub use taffy::style::AvailableSpace;
-use taffy::{node::MeasureFunc, prelude::Size as TaffySize};
+
+use crate::widget::ImageMeasure;
+
+#[cfg(feature = "bevy_text")]
+use crate::widget::TextMeasure;
 
 impl std::fmt::Debug for ContentSize {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
@@ -25,6 +29,44 @@ pub trait Measure: Send + Sync + 'static {
     ) -> Vec2;
 }
 
+/// A type to serve as Taffy's node context (which allows the content size of leaf nodes to be computed)
+///
+/// It has specific variants for common built-in types to avoid making them opaque and needing to box them
+/// by wrapping them in a closure and a Custom variant that allows arbitrary measurement closures if required.
+pub enum NodeMeasure {
+    Fixed(FixedMeasure),
+    #[cfg(feature = "bevy_text")]
+    Text(TextMeasure),
+    Image(ImageMeasure),
+    Custom(Box<dyn Measure>),
+}
+
+impl Measure for NodeMeasure {
+    fn measure(
+        &self,
+        width: Option<f32>,
+        height: Option<f32>,
+        available_width: AvailableSpace,
+        available_height: AvailableSpace,
+    ) -> Vec2 {
+        match self {
+            NodeMeasure::Fixed(fixed) => {
+                fixed.measure(width, height, available_width, available_height)
+            }
+            #[cfg(feature = "bevy_text")]
+            NodeMeasure::Text(text) => {
+                text.measure(width, height, available_width, available_height)
+            }
+            NodeMeasure::Image(image) => {
+                image.measure(width, height, available_width, available_height)
+            }
+            NodeMeasure::Custom(custom) => {
+                custom.measure(width, height, available_width, available_height)
+            }
+        }
+    }
+}
+
 /// A `FixedMeasure` is a `Measure` that ignores all constraints and
 /// always returns the same size.
 #[derive(Default, Clone)]
@@ -51,26 +93,19 @@ impl Measure for FixedMeasure {
 pub struct ContentSize {
     /// The `Measure` used to compute the intrinsic size
     #[reflect(ignore)]
-    pub(crate) measure_func: Option<MeasureFunc>,
+    pub(crate) measure: Option<NodeMeasure>,
 }
 
 impl ContentSize {
     /// Set a `Measure` for the UI node entity with this component
-    pub fn set(&mut self, measure: impl Measure) {
-        let measure_func = move |size: TaffySize<_>, available: TaffySize<_>| {
-            let size = measure.measure(size.width, size.height, available.width, available.height);
-            TaffySize {
-                width: size.x,
-                height: size.y,
-            }
-        };
-        self.measure_func = Some(MeasureFunc::Boxed(Box::new(measure_func)));
+    pub fn set(&mut self, measure: NodeMeasure) {
+        self.measure = Some(measure);
     }
 
     /// Creates a `ContentSize` with a `Measure` that always returns given `size` argument, regardless of the UI layout's constraints.
     pub fn fixed_size(size: Vec2) -> ContentSize {
         let mut content_size = Self::default();
-        content_size.set(FixedMeasure { size });
+        content_size.set(NodeMeasure::Fixed(FixedMeasure { size }));
         content_size
     }
 }
diff --git a/crates/bevy_ui/src/ui_node.rs b/crates/bevy_ui/src/ui_node.rs
index 00d59bcc6f27f..75852e78ee212 100644
--- a/crates/bevy_ui/src/ui_node.rs
+++ b/crates/bevy_ui/src/ui_node.rs
@@ -770,10 +770,12 @@ impl Default for Direction {
     reflect(Serialize, Deserialize)
 )]
 pub enum Display {
-    /// Use Flexbox layout model to determine the position of this [`Node`].
+    /// Use Flexbox layout model to determine the position of this [`Node`]'s children.
     Flex,
-    /// Use CSS Grid layout model to determine the position of this [`Node`].
+    /// Use CSS Grid layout model to determine the position of this [`Node`]'s children.
     Grid,
+    /// Use CSS Block layout model to determine the position of this [`Node`]'s children.
+    Block,
     /// Use no layout, don't render this node and its children.
     ///
     /// If you want to hide a node and its children,
@@ -896,8 +898,10 @@ impl Default for Overflow {
 pub enum OverflowAxis {
     /// Show overflowing items.
     Visible,
-    /// Hide overflowing items.
+    /// Hide overflowing items by clipping.
     Clip,
+    /// Hide overflowing items by influencing layout and then clipping.
+    Hidden,
 }
 
 impl OverflowAxis {
diff --git a/crates/bevy_ui/src/widget/image.rs b/crates/bevy_ui/src/widget/image.rs
index 4c49dd1af9e21..b78515c68d2ff 100644
--- a/crates/bevy_ui/src/widget/image.rs
+++ b/crates/bevy_ui/src/widget/image.rs
@@ -1,4 +1,6 @@
-use crate::{measurement::AvailableSpace, ContentSize, Measure, Node, UiImage, UiScale};
+use crate::{
+    measurement::AvailableSpace, ContentSize, Measure, Node, NodeMeasure, UiImage, UiScale,
+};
 use bevy_asset::Assets;
 use bevy_ecs::prelude::*;
 use bevy_math::{UVec2, Vec2};
@@ -72,6 +74,7 @@ pub fn update_image_content_size_system(
     windows: Query<&Window, With<PrimaryWindow>>,
     ui_scale: Res<UiScale>,
     textures: Res<Assets<Image>>,
+
     atlases: Res<Assets<TextureAtlasLayout>>,
     mut query: Query<
         (
@@ -100,10 +103,10 @@ pub fn update_image_content_size_system(
                 || content_size.is_added()
             {
                 image_size.size = size;
-                content_size.set(ImageMeasure {
+                content_size.set(NodeMeasure::Image(ImageMeasure {
                     // multiply the image size by the scale factor to get the physical size
                     size: size.as_vec2() * combined_scale_factor,
-                });
+                }));
             }
         }
     }
diff --git a/crates/bevy_ui/src/widget/text.rs b/crates/bevy_ui/src/widget/text.rs
index 09ee5049781fe..5aa6b44ed6f57 100644
--- a/crates/bevy_ui/src/widget/text.rs
+++ b/crates/bevy_ui/src/widget/text.rs
@@ -1,4 +1,4 @@
-use crate::{ContentSize, FixedMeasure, Measure, Node, UiScale};
+use crate::{ContentSize, FixedMeasure, Measure, Node, NodeMeasure, UiScale};
 use bevy_asset::Assets;
 use bevy_ecs::{
     prelude::{Component, DetectChanges},
@@ -88,9 +88,9 @@ fn create_text_measure(
     match TextMeasureInfo::from_text(&text, fonts, scale_factor) {
         Ok(measure) => {
             if text.linebreak_behavior == BreakLineOn::NoWrap {
-                content_size.set(FixedMeasure { size: measure.max });
+                content_size.set(NodeMeasure::Fixed(FixedMeasure { size: measure.max }));
             } else {
-                content_size.set(TextMeasure { info: measure });
+                content_size.set(NodeMeasure::Text(TextMeasure { info: measure }));
             }
 
             // Text measure func created successfully, so set `TextFlags` to schedule a recompute
diff --git a/examples/ui/display_and_visibility.rs b/examples/ui/display_and_visibility.rs
index d95e7b8c56cc2..6c6b4112487fd 100644
--- a/examples/ui/display_and_visibility.rs
+++ b/examples/ui/display_and_visibility.rs
@@ -54,7 +54,7 @@ impl TargetUpdate for Target<Display> {
         style.display = match style.display {
             Display::Flex => Display::None,
             Display::None => Display::Flex,
-            Display::Grid => unreachable!(),
+            Display::Block | Display::Grid => unreachable!(),
         };
         format!("{}::{:?} ", Self::NAME, style.display)
     }