Add context type parameter to Taffy struct

src/compute/leaf.rs

@@ -1,6 +1,6 @@
 //! Computes size using styles and measure functions
 
-use crate::geometry::{Line, Point, Size};
+use crate::geometry::{Point, Size};
 use crate::style::{AvailableSpace, Display, Overflow, Position, Style};
 use crate::tree::{CollapsibleMarginSet, Measurable};
 use crate::tree::{SizeBaselinesAndMargins, SizingMode};
@@ -12,39 +12,40 @@ use crate::util::{MaybeResolve, ResolveOrZero};
 use crate::util::debug::NODE_LOGGER;
 
 /// Perform full layout on a leaf node
-pub(crate) fn perform_layout(
+pub(crate) fn perform_layout<Context>(
     style: &Style,
-    measurable: Option<&impl Measurable>,
+    measurable: Option<&impl Measurable<Context = Context>>,
     known_dimensions: Size<Option<f32>>,
     parent_size: Size<Option<f32>>,
     available_space: Size<AvailableSpace>,
     sizing_mode: SizingMode,
-    _vertical_margins_are_collapsible: Line<bool>,
+    context: &mut Context,
 ) -> SizeBaselinesAndMargins {
-    compute(style, measurable, known_dimensions, parent_size, available_space, sizing_mode)
+    compute(style, measurable, known_dimensions, parent_size, available_space, sizing_mode, context)
 }
 
 /// Measure a leaf node's size
-pub(crate) fn measure_size(
+pub(crate) fn measure_size<Context>(
     style: &Style,
-    measurable: Option<&impl Measurable>,
+    measurable: Option<&impl Measurable<Context = Context>>,
     known_dimensions: Size<Option<f32>>,
     parent_size: Size<Option<f32>>,
     available_space: Size<AvailableSpace>,
     sizing_mode: SizingMode,
-    _vertical_margins_are_collapsible: Line<bool>,
+    context: &mut Context,
 ) -> Size<f32> {
-    compute(style, measurable, known_dimensions, parent_size, available_space, sizing_mode).size
+    compute(style, measurable, known_dimensions, parent_size, available_space, sizing_mode, context).size
 }
 
 /// Compute the size of a leaf node (node with no children)
-pub fn compute(
+pub fn compute<Context>(
     style: &Style,
-    measurable: Option<&impl Measurable>,
+    measurable: Option<&impl Measurable<Context = Context>>,
     known_dimensions: Size<Option<f32>>,
     parent_size: Size<Option<f32>>,
     available_space: Size<AvailableSpace>,
     sizing_mode: SizingMode,
+    context: &mut Context,
 ) -> SizeBaselinesAndMargins {
     // Resolve node's preferred/min/max sizes (width/heights) against the available space (percentages resolve to pixel values)
     // For ContentSize mode, we pretend that the node has no size styles as these should be ignored.
@@ -141,7 +142,7 @@ pub fn compute(
         };
 
         // Measure node
-        let measured_size = measurable.measure(known_dimensions, available_space);
+        let measured_size = measurable.measure(known_dimensions, available_space, context);
         let clamped_size = node_size.unwrap_or(measured_size).maybe_clamp(node_min_size, node_max_size);
         let size = Size {
             width: clamped_size.width,

src/compute/taffy_tree.rs

@@ -3,7 +3,9 @@
 use crate::compute::{leaf, LayoutAlgorithm};
 use crate::geometry::{Line, Point, Size};
 use crate::style::{AvailableSpace, Display};
-use crate::tree::{Layout, LayoutTree, NodeId, RunMode, SizeBaselinesAndMargins, SizingMode, Taffy, TaffyError};
+use crate::tree::{
+    Layout, LayoutTree, Measurable, NodeId, RunMode, SizeBaselinesAndMargins, SizingMode, Taffy, TaffyError,
+};
 use crate::util::sys::round;
 
 #[cfg(feature = "block_layout")]
@@ -34,8 +36,8 @@ fn debug_log_node(
 }
 
 /// Updates the stored layout of the provided `node` and its children
-pub(crate) fn compute_layout(
-    taffy: &mut Taffy,
+pub(crate) fn compute_layout<Measure: Measurable>(
+    taffy: &mut Taffy<Measure>,
     root: NodeId,
     available_space: Size<AvailableSpace>,
 ) -> Result<(), TaffyError> {
@@ -62,8 +64,8 @@ pub(crate) fn compute_layout(
 }
 
 /// Perform full layout on a node. Chooses which algorithm to use based on the `display` property.
-pub(crate) fn perform_node_layout(
-    tree: &mut Taffy,
+pub(crate) fn perform_node_layout<Measure: Measurable>(
+    tree: &mut Taffy<Measure>,
     node: NodeId,
     known_dimensions: Size<Option<f32>>,
     parent_size: Size<Option<f32>>,
@@ -84,8 +86,8 @@ pub(crate) fn perform_node_layout(
 }
 
 /// Measure a node's size. Chooses which algorithm to use based on the `display` property.
-pub(crate) fn measure_node_size(
-    tree: &mut Taffy,
+pub(crate) fn measure_node_size<Measure: Measurable>(
+    tree: &mut Taffy<Measure>,
     node: NodeId,
     known_dimensions: Size<Option<f32>>,
     parent_size: Size<Option<f32>>,
@@ -108,8 +110,8 @@ pub(crate) fn measure_node_size(
 
 /// Updates the stored layout of the provided `node` and its children
 #[allow(clippy::too_many_arguments)]
-fn compute_node_layout(
-    tree: &mut Taffy,
+fn compute_node_layout<Measure: Measurable>(
+    tree: &mut Taffy<Measure>,
     node: NodeId,
     known_dimensions: Size<Option<f32>>,
     parent_size: Size<Option<f32>>,
@@ -228,7 +230,7 @@ fn compute_node_layout(
                 parent_size,
                 available_space,
                 sizing_mode,
-                vertical_margins_are_collapsible,
+                tree.context.as_mut().unwrap(),
             ),
             RunMode::ComputeSize => leaf::measure_size(
                 &tree.nodes[node_key].style,
@@ -237,7 +239,7 @@ fn compute_node_layout(
                 parent_size,
                 available_space,
                 sizing_mode,
-                vertical_margins_are_collapsible,
+                tree.context.as_mut().unwrap(),
             )
             .into(),
         },
@@ -256,9 +258,9 @@ fn compute_node_layout(
 
 /// Creates a layout for this node and its children, recursively.
 /// Each hidden node has zero size and is placed at the origin
-fn perform_taffy_tree_hidden_layout(tree: &mut Taffy, node: NodeId) {
+fn perform_taffy_tree_hidden_layout<Measure: Measurable>(tree: &mut Taffy<Measure>, node: NodeId) {
     /// Recursive function to apply hidden layout to all descendents
-    fn perform_hidden_layout_inner(tree: &mut Taffy, node: NodeId, order: u32) {
+    fn perform_hidden_layout_inner<Measure: Measurable>(tree: &mut Taffy<Measure>, node: NodeId, order: u32) {
         let node_key = node.into();
         *tree.layout_mut(node) = Layout::with_order(order);
         tree.nodes[node_key].cache.clear();

src/tree/measure_func.rs

@@ -18,14 +18,17 @@ pub trait Measurable: Send + Sync {
         &self,
         known_dimensions: Size<Option<f32>>,
         available_space: Size<AvailableSpace>,
-        context: Self::Context,
+        context: &mut Self::Context,
     ) -> Size<f32>;
 }
 
 /// A function that can be used to compute the intrinsic size of a node
 pub enum MeasureFunc<Context = ()> {
-    /// Stores an unboxed function
-    Raw(fn(Size<Option<f32>>, Size<AvailableSpace>, context: Context) -> Size<f32>),
+    /// Stores an unboxed function with no context parameter
+    Raw(fn(Size<Option<f32>>, Size<AvailableSpace>) -> Size<f32>),
+
+    /// Stores an unboxed function with a context parameter
+    RawWithContext(fn(Size<Option<f32>>, Size<AvailableSpace>, context: &mut Context) -> Size<f32>),
 
     /// Stores a boxed function
     #[cfg(any(feature = "std", feature = "alloc"))]
@@ -37,9 +40,15 @@ impl<Context> Measurable for MeasureFunc<Context> {
 
     /// Call the measure function to measure to the node
     #[inline(always)]
-    fn measure(&self, known_dimensions: Size<Option<f32>>, available_space: Size<AvailableSpace>, context: Context) -> Size<f32> {
+    fn measure(
+        &self,
+        known_dimensions: Size<Option<f32>>,
+        available_space: Size<AvailableSpace>,
+        context: &mut Context,
+    ) -> Size<f32> {
         match self {
-            Self::Raw(measure) => measure(known_dimensions, available_space, context),
+            Self::Raw(measure) => measure(known_dimensions, available_space),
+            Self::RawWithContext(measure) => measure(known_dimensions, available_space, context),
             #[cfg(any(feature = "std", feature = "alloc"))]
             Self::Boxed(measurable) => measurable.measure(known_dimensions, available_space, context),
         }

src/tree/taffy_tree/tree.rs

@@ -7,7 +7,7 @@ use crate::compute::taffy_tree::{compute_layout, measure_node_size, perform_node
 use crate::geometry::{Line, Size};
 use crate::prelude::LayoutTree;
 use crate::style::{AvailableSpace, Style};
-use crate::tree::{Layout, MeasureFunc, NodeData, NodeId, SizeBaselinesAndMargins, SizingMode};
+use crate::tree::{Layout, Measurable, MeasureFunc, NodeData, NodeId, SizeBaselinesAndMargins, SizingMode};
 use crate::util::sys::{new_vec_with_capacity, ChildrenVec, Vec};
 
 use super::{TaffyError, TaffyResult};
@@ -25,12 +25,15 @@ impl Default for TaffyConfig {
 }
 
 /// A tree of UI nodes suitable for UI layout
-pub struct Taffy {
+pub struct Taffy<Measure = MeasureFunc<()>>
+where
+    Measure: Measurable,
+{
     /// The [`NodeData`] for each node stored in this tree
     pub(crate) nodes: SlotMap<DefaultKey, NodeData>,
 
     /// Functions/closures that compute the intrinsic size of leaf nodes
-    pub(crate) measure_funcs: SparseSecondaryMap<DefaultKey, MeasureFunc>,
+    pub(crate) measure_funcs: SparseSecondaryMap<DefaultKey, MeasureFunc<Measure::Context>>,
 
     /// The children of each node
     ///
@@ -44,6 +47,9 @@ pub struct Taffy {
 
     /// Layout mode configuration
     pub(crate) config: TaffyConfig,
+
+    /// Used to store the context during layout. Cleared before returning from `compute_layout`.
+    pub(crate) context: Option<Measure::Context>,
 }
 
 impl Default for Taffy {
@@ -62,8 +68,8 @@ impl<'a> Iterator for TaffyChildIter<'a> {
     }
 }
 
-impl LayoutTree for Taffy {
-    type ChildIter<'a> = TaffyChildIter<'a>;
+impl<Measure: Measurable> LayoutTree for Taffy<Measure> {
+    type ChildIter<'a> = TaffyChildIter<'a> where Measure: 'a;
 
     #[inline(always)]
     fn children(&self, node: NodeId) -> Self::ChildIter<'_> {
@@ -139,29 +145,7 @@ impl LayoutTree for Taffy {
 }
 
 #[allow(clippy::iter_cloned_collect)] // due to no-std support, we need to use `iter_cloned` instead of `collect`
-impl Taffy {
-    /// Creates a new [`Taffy`]
-    ///
-    /// The default capacity of a [`Taffy`] is 16 nodes.
-    #[must_use]
-    pub fn new() -> Self {
-        Self::with_capacity(16)
-    }
-
-    /// Creates a new [`Taffy`] that can store `capacity` nodes before reallocation
-    #[must_use]
-    pub fn with_capacity(capacity: usize) -> Self {
-        Self {
-            // TODO: make this method const upstream,
-            // so constructors here can be const
-            nodes: SlotMap::with_capacity(capacity),
-            children: SlotMap::with_capacity(capacity),
-            parents: SlotMap::with_capacity(capacity),
-            measure_funcs: SparseSecondaryMap::with_capacity(capacity),
-            config: TaffyConfig::default(),
-        }
-    }
-
+impl<Measure: Measurable> Taffy<Measure> {
     /// Enable rounding of layout values. Rounding is enabled by default.
     pub fn enable_rounding(&mut self) {
         self.config.use_rounding = true;
@@ -184,7 +168,11 @@ impl Taffy {
     /// Creates and adds a new unattached leaf node to the tree, and returns the node of the new node
     ///
     /// Creates and adds a new leaf node with a supplied [`MeasureFunc`]
-    pub fn new_leaf_with_measure(&mut self, layout: Style, measure: MeasureFunc) -> TaffyResult<NodeId> {
+    pub fn new_leaf_with_measure(
+        &mut self,
+        layout: Style,
+        measure: MeasureFunc<Measure::Context>,
+    ) -> TaffyResult<NodeId> {
         let mut data = NodeData::new(layout);
         data.needs_measure = true;
 
@@ -237,7 +225,7 @@ impl Taffy {
     }
 
     /// Sets the [`MeasureFunc`] of the associated node
-    pub fn set_measure(&mut self, node: NodeId, measure: Option<MeasureFunc>) -> TaffyResult<()> {
+    pub fn set_measure(&mut self, node: NodeId, measure: Option<MeasureFunc<Measure::Context>>) -> TaffyResult<()> {
         let key = node.into();
         if let Some(measure) = measure {
             self.nodes[key].needs_measure = true;
@@ -404,9 +392,47 @@ impl Taffy {
         Ok(self.nodes[node.into()].cache.is_empty())
     }
 
+    /// Updates the stored layout of the provided `node` and its children
+    pub fn compute_layout_with_context(
+        &mut self,
+        node: NodeId,
+        available_space: Size<AvailableSpace>,
+        context: Measure::Context,
+    ) -> Result<(), TaffyError> {
+        self.context = Some(context);
+        let result = compute_layout(self, node, available_space);
+        self.context = None;
+        result
+    }
+}
+
+impl Taffy<MeasureFunc<()>> {
+    /// Creates a new [`Taffy`]
+    ///
+    /// The default capacity of a [`Taffy`] is 16 nodes.
+    #[must_use]
+    pub fn new() -> Self {
+        Self::with_capacity(16)
+    }
+
+    /// Creates a new [`Taffy`] that can store `capacity` nodes before reallocation
+    #[must_use]
+    pub fn with_capacity(capacity: usize) -> Self {
+        Taffy {
+            // TODO: make this method const upstream,
+            // so constructors here can be const
+            nodes: SlotMap::with_capacity(capacity),
+            children: SlotMap::with_capacity(capacity),
+            parents: SlotMap::with_capacity(capacity),
+            measure_funcs: SparseSecondaryMap::with_capacity(capacity),
+            config: TaffyConfig::default(),
+            context: None,
+        }
+    }
+
     /// Updates the stored layout of the provided `node` and its children
     pub fn compute_layout(&mut self, node: NodeId, available_space: Size<AvailableSpace>) -> Result<(), TaffyError> {
-        compute_layout(self, node, available_space)
+        self.compute_layout_with_context(node, available_space, ())
     }
 }