From 3b0815d618374510f8a16390990d29107ed947ae Mon Sep 17 00:00:00 2001
From: Nico Burns <nico@nicoburns.com>
Date: Thu, 24 Aug 2023 01:13:52 +0100
Subject: [PATCH] Implement TaffyView

---
 src/compute/mod.rs          |  6 ++--
 src/compute/taffy_tree.rs   | 63 ++++++++++++++++++-------------------
 src/tree/measure_func.rs    | 15 +++++++++
 src/tree/mod.rs             |  2 +-
 src/tree/taffy_tree/mod.rs  |  2 +-
 src/tree/taffy_tree/tree.rs | 54 ++++++++++++++++---------------
 src/util/debug.rs           | 17 +++++-----
 7 files changed, 88 insertions(+), 71 deletions(-)

diff --git a/src/compute/mod.rs b/src/compute/mod.rs
index 8f61d8140..3537647af 100644
--- a/src/compute/mod.rs
+++ b/src/compute/mod.rs
@@ -110,12 +110,12 @@ mod tests {
     use super::perform_hidden_layout;
     use crate::geometry::{Point, Size};
     use crate::style::{Display, Style};
-    use crate::tree::MeasureFunc;
+    use crate::tree::TaffyView;
     use crate::Taffy;
 
     #[test]
     fn hidden_layout_should_hide_recursively() {
-        let mut taffy: Taffy<MeasureFunc<()>> = Taffy::new();
+        let mut taffy: Taffy<()> = Taffy::new();
 
         let style: Style = Style { display: Display::Flex, size: Size::from_lengths(50.0, 50.0), ..Default::default() };
 
@@ -133,7 +133,7 @@ mod tests {
             )
             .unwrap();
 
-        perform_hidden_layout(&mut taffy, root.into());
+        perform_hidden_layout(&mut TaffyView { taffy: &mut taffy, context: () }, root.into());
 
         // Whatever size and display-mode the nodes had previously,
         // all layouts should resolve to ZERO due to the root's DISPLAY::NONE
diff --git a/src/compute/taffy_tree.rs b/src/compute/taffy_tree.rs
index 77f56d099..a5e4905ea 100644
--- a/src/compute/taffy_tree.rs
+++ b/src/compute/taffy_tree.rs
@@ -4,7 +4,7 @@ use crate::compute::{leaf, LayoutAlgorithm};
 use crate::geometry::{Line, Point, Size};
 use crate::style::{AvailableSpace, Display};
 use crate::tree::{
-    Layout, LayoutTree, Measurable, NodeId, RunMode, SizeBaselinesAndMargins, SizingMode, Taffy, TaffyError,
+    Layout, LayoutTree, Measurable, NodeId, RunMode, SizeBaselinesAndMargins, SizingMode, Taffy, TaffyView, TaffyError,
 };
 use crate::util::sys::round;
 
@@ -37,15 +37,14 @@ fn debug_log_node(
 
 /// Updates the stored layout of the provided `node` and its children
 pub(crate) fn compute_layout<Measure: Measurable>(
-    taffy: &mut Taffy<Measure>,
+    taffy_view: &mut TaffyView<Measure>,
     root: NodeId,
     available_space: Size<AvailableSpace>,
 ) -> Result<(), TaffyError> {
-    taffy.is_layouting = true;
 
     // Recursively compute node layout
     let size_and_baselines = perform_node_layout(
-        taffy,
+        taffy_view,
         root,
         Size::NONE,
         available_space.into_options(),
@@ -55,21 +54,19 @@ pub(crate) fn compute_layout<Measure: Measurable>(
     );
 
     let layout = Layout { order: 0, size: size_and_baselines.size, location: Point::ZERO };
-    *taffy.layout_mut(root) = layout;
+    *taffy_view.layout_mut(root) = layout;
 
     // If rounding is enabled, recursively round the layout's of this node and all children
-    if taffy.config.use_rounding {
-        round_layout(taffy, root, 0.0, 0.0);
+    if taffy_view.taffy.config.use_rounding {
+        round_layout(taffy_view.taffy, root, 0.0, 0.0);
     }
 
-    taffy.is_layouting = false;
-
     Ok(())
 }
 
 /// Perform full layout on a node. Chooses which algorithm to use based on the `display` property.
 pub(crate) fn perform_node_layout<Measure: Measurable>(
-    tree: &mut Taffy<Measure>,
+    taffy_view: &mut TaffyView<Measure>,
     node: NodeId,
     known_dimensions: Size<Option<f32>>,
     parent_size: Size<Option<f32>>,
@@ -78,7 +75,7 @@ pub(crate) fn perform_node_layout<Measure: Measurable>(
     vertical_margins_are_collapsible: Line<bool>,
 ) -> SizeBaselinesAndMargins {
     compute_node_layout(
-        tree,
+        taffy_view,
         node,
         known_dimensions,
         parent_size,
@@ -91,7 +88,7 @@ pub(crate) fn perform_node_layout<Measure: Measurable>(
 
 /// Measure a node's size. Chooses which algorithm to use based on the `display` property.
 pub(crate) fn measure_node_size<Measure: Measurable>(
-    tree: &mut Taffy<Measure>,
+    taffy_view: &mut TaffyView<Measure>,
     node: NodeId,
     known_dimensions: Size<Option<f32>>,
     parent_size: Size<Option<f32>>,
@@ -100,7 +97,7 @@ pub(crate) fn measure_node_size<Measure: Measurable>(
     vertical_margins_are_collapsible: Line<bool>,
 ) -> Size<f32> {
     compute_node_layout(
-        tree,
+        taffy_view,
         node,
         known_dimensions,
         parent_size,
@@ -115,7 +112,7 @@ pub(crate) fn measure_node_size<Measure: Measurable>(
 /// Updates the stored layout of the provided `node` and its children
 #[allow(clippy::too_many_arguments)]
 fn compute_node_layout<Measure: Measurable>(
-    tree: &mut Taffy<Measure>,
+    taffy_view: &mut TaffyView<Measure>,
     node: NodeId,
     known_dimensions: Size<Option<f32>>,
     parent_size: Size<Option<f32>>,
@@ -130,12 +127,12 @@ fn compute_node_layout<Measure: Measurable>(
     println!();
 
     let node_key = node.into();
-    let has_children = !tree.children[node_key].is_empty();
+    let has_children = !taffy_view.taffy.children[node_key].is_empty();
 
     // First we check if we have a cached result for the given input
     let cache_run_mode = if !has_children { RunMode::PerformLayout } else { run_mode };
     if let Some(cached_size_and_baselines) =
-        tree.nodes[node_key].cache.get(known_dimensions, available_space, cache_run_mode)
+        taffy_view.taffy.nodes[node_key].cache.get(known_dimensions, available_space, cache_run_mode)
     {
         #[cfg(feature = "debug")]
         NODE_LOGGER.labelled_debug_log("CACHE", cached_size_and_baselines.size);
@@ -188,15 +185,15 @@ fn compute_node_layout<Measure: Measurable>(
         }
     }
 
-    let display_mode = tree.nodes[node_key].style.display;
+    let display_mode = taffy_view.taffy.nodes[node_key].style.display;
     let computed_size_and_baselines = match (display_mode, has_children) {
         (Display::None, _) => {
-            perform_taffy_tree_hidden_layout(tree, node);
+            perform_taffy_tree_hidden_layout(taffy_view.taffy, node);
             SizeBaselinesAndMargins::HIDDEN
         }
         #[cfg(feature = "block_layout")]
         (Display::Block, true) => perform_computations::<BlockAlgorithm>(
-            tree,
+            taffy_view,
             node,
             known_dimensions,
             parent_size,
@@ -207,7 +204,7 @@ fn compute_node_layout<Measure: Measurable>(
         ),
         #[cfg(feature = "flexbox")]
         (Display::Flex, true) => perform_computations::<FlexboxAlgorithm>(
-            tree,
+            taffy_view,
             node,
             known_dimensions,
             parent_size,
@@ -218,7 +215,7 @@ fn compute_node_layout<Measure: Measurable>(
         ),
         #[cfg(feature = "grid")]
         (Display::Grid, true) => perform_computations::<CssGridAlgorithm>(
-            tree,
+            taffy_view,
             node,
             known_dimensions,
             parent_size,
@@ -229,29 +226,29 @@ fn compute_node_layout<Measure: Measurable>(
         ),
         (_, false) => match run_mode {
             RunMode::PerformLayout => leaf::perform_layout(
-                &tree.nodes[node_key].style,
-                tree.nodes[node_key].needs_measure.then(|| &mut tree.measure_funcs[node_key]),
+                &taffy_view.taffy.nodes[node_key].style,
+                taffy_view.taffy.nodes[node_key].needs_measure.then(|| &mut taffy_view.taffy.measure_funcs[node_key]),
                 known_dimensions,
                 parent_size,
                 available_space,
                 sizing_mode,
-                tree.context.as_mut().unwrap(),
+                &mut taffy_view.context,
             ),
             RunMode::ComputeSize => leaf::measure_size(
-                &tree.nodes[node_key].style,
-                tree.nodes[node_key].needs_measure.then(|| &mut tree.measure_funcs[node_key]),
+                &taffy_view.taffy.nodes[node_key].style,
+                taffy_view.taffy.nodes[node_key].needs_measure.then(|| &mut taffy_view.taffy.measure_funcs[node_key]),
                 known_dimensions,
                 parent_size,
                 available_space,
                 sizing_mode,
-                tree.context.as_mut().unwrap(),
+                &mut taffy_view.context,
             )
             .into(),
         },
     };
 
     // Cache result
-    tree.nodes[node_key].cache.store(known_dimensions, available_space, cache_run_mode, computed_size_and_baselines);
+    taffy_view.taffy.nodes[node_key].cache.store(known_dimensions, available_space, cache_run_mode, computed_size_and_baselines);
 
     #[cfg(feature = "debug")]
     NODE_LOGGER.labelled_debug_log("RESULT", computed_size_and_baselines.size);
@@ -267,15 +264,17 @@ fn perform_taffy_tree_hidden_layout<Measure: Measurable>(tree: &mut Taffy<Measur
     /// Recursive function to apply hidden layout to all descendents
     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].unrounded_layout = Layout::with_order(order);
+        tree.nodes[node_key].final_layout = Layout::with_order(order);
         tree.nodes[node_key].cache.clear();
         for order in 0..tree.children[node_key].len() {
-            perform_hidden_layout_inner(tree, tree.child(node, order), order as _);
+            perform_hidden_layout_inner(tree, tree.children[node_key][order], order as _);
         }
     }
 
+    let node_key = node.into();
     for order in 0..tree.children[node.into()].len() {
-        perform_hidden_layout_inner(tree, tree.child(node, order), order as _);
+        perform_hidden_layout_inner(tree, tree.children[node_key][order], order as _);
     }
 }
 
@@ -305,7 +304,7 @@ fn round_layout<Measure: Measurable>(tree: &mut Taffy<Measure>, node_id: NodeId,
 
     let child_count = tree.child_count(node_id).unwrap();
     for index in 0..child_count {
-        let child = tree.child(node_id, index);
+        let child = tree.children[node_id.into()][index];
         round_layout(tree, child, cumulative_x, cumulative_y);
     }
 }
diff --git a/src/tree/measure_func.rs b/src/tree/measure_func.rs
index 6e0e4dd1d..62f731f5d 100644
--- a/src/tree/measure_func.rs
+++ b/src/tree/measure_func.rs
@@ -56,6 +56,21 @@ impl<Context> MeasureFunc<Context> {
     }
 }
 
+impl Measurable for () {
+    type Context = ();
+
+    /// Call the measure function to measure the node
+    #[inline(always)]
+    fn measure(
+        &mut self,
+        _known_dimensions: Size<Option<f32>>,
+        _available_space: Size<AvailableSpace>,
+        _context: &mut (),
+    ) -> Size<f32> {
+        Size::ZERO
+    }
+}
+
 impl<Context> Measurable for MeasureFunc<Context> {
     type Context = Context;
 
diff --git a/src/tree/mod.rs b/src/tree/mod.rs
index efacdb99d..e2154b0a4 100644
--- a/src/tree/mod.rs
+++ b/src/tree/mod.rs
@@ -15,7 +15,7 @@ pub use node::NodeId;
 #[cfg(feature = "taffy_tree")]
 mod taffy_tree;
 #[cfg(feature = "taffy_tree")]
-pub use taffy_tree::{Taffy, TaffyChildIter, TaffyError, TaffyResult};
+pub use taffy_tree::{Taffy, TaffyView, TaffyChildIter, TaffyError, TaffyResult};
 mod layout;
 pub use layout::{CollapsibleMarginSet, Layout, RunMode, SizeBaselinesAndMargins, SizingMode};
 
diff --git a/src/tree/taffy_tree/mod.rs b/src/tree/taffy_tree/mod.rs
index 64f7fce4e..596089c7b 100644
--- a/src/tree/taffy_tree/mod.rs
+++ b/src/tree/taffy_tree/mod.rs
@@ -4,4 +4,4 @@ mod error;
 mod tree;
 
 pub use error::{TaffyError, TaffyResult};
-pub use tree::{Taffy, TaffyChildIter};
+pub use tree::{Taffy, TaffyView, TaffyChildIter};
diff --git a/src/tree/taffy_tree/tree.rs b/src/tree/taffy_tree/tree.rs
index 79adf1c05..185d16601 100644
--- a/src/tree/taffy_tree/tree.rs
+++ b/src/tree/taffy_tree/tree.rs
@@ -47,14 +47,6 @@ where
 
     /// Layout mode configuration
     pub(crate) config: TaffyConfig,
-
-    /// Hack to allow the `LayoutTree::layout_mut` function to expose the `NodeData.unrounded_layout` of a node to
-    /// the layout algorithms during layout, while exposing the `NodeData.final_layout` when called by external users.
-    /// This allows us to fix <https://github.com/DioxusLabs/taffy/issues/501> without breaking backwards compatibility
-    pub(crate) is_layouting: bool,
-
-    /// Used to store the context during layout. Cleared before returning from `compute_layout`.
-    pub(crate) context: Option<Measure::Context>,
 }
 
 impl Default for Taffy {
@@ -73,45 +65,56 @@ impl<'a> Iterator for TaffyChildIter<'a> {
     }
 }
 
-impl<Measure: Measurable> LayoutTree for Taffy<Measure> {
-    type ChildIter<'a> = TaffyChildIter<'a> where Measure: 'a;
+pub struct TaffyView<'t, Measure: Measurable> {
+    pub (crate) taffy: &'t mut Taffy<Measure>,
+    pub (crate) context: Measure::Context,
+}
+
+impl <'t, Measure: Measurable> TaffyView<'t, Measure> {
+    pub (crate) fn config(&self) -> &TaffyConfig {
+        &self.taffy.config
+    }
+}
+
+impl <'t, Measure: Measurable> LayoutTree for TaffyView<'t, Measure> {
+    type ChildIter<'a> = TaffyChildIter<'a> where Measure: 'a, Self: 'a;
 
     #[inline(always)]
     fn children(&self, node: NodeId) -> Self::ChildIter<'_> {
-        TaffyChildIter(self.children[node.into()].iter())
+        TaffyChildIter(self.taffy.children[node.into()].iter())
     }
 
     #[inline(always)]
     fn child_count(&self, node: NodeId) -> usize {
-        self.children[node.into()].len()
+        self.taffy.children[node.into()].len()
     }
 
     #[inline(always)]
     fn style(&self, node: NodeId) -> &Style {
-        &self.nodes[node.into()].style
+        &self.taffy.nodes[node.into()].style
     }
 
     #[inline(always)]
     fn layout(&self, node: NodeId) -> &Layout {
-        if self.is_layouting && self.config.use_rounding {
-            &self.nodes[node.into()].unrounded_layout
+        if self.taffy.config.use_rounding {
+            &self.taffy.nodes[node.into()].unrounded_layout
         } else {
-            &self.nodes[node.into()].final_layout
+            &self.taffy.nodes[node.into()].final_layout
         }
     }
 
     #[inline(always)]
     fn layout_mut(&mut self, node: NodeId) -> &mut Layout {
-        if self.is_layouting && self.config.use_rounding {
-            &mut self.nodes[node.into()].unrounded_layout
+        if self.taffy.config.use_rounding {
+            &mut self.taffy.nodes[node.into()].unrounded_layout
         } else {
-            &mut self.nodes[node.into()].final_layout
+            &mut self.taffy.nodes[node.into()].final_layout
         }
     }
 
     #[inline(always)]
     fn child(&self, node: NodeId, id: usize) -> NodeId {
-        self.children[node.into()][id]
+        self.taffy.children[node.into()][id]
     }
 
     #[inline(always)]
@@ -178,8 +181,6 @@ impl<Measure: Measurable> Taffy<Measure> {
             parents: SlotMap::with_capacity(capacity),
             measure_funcs: SparseSecondaryMap::with_capacity(capacity),
             config: TaffyConfig::default(),
-            is_layouting: false,
-            context: None,
         }
     }
 
@@ -455,10 +456,11 @@ impl<Measure: Measurable> Taffy<Measure> {
         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
+        let mut taffy_view = TaffyView {
+            taffy: self,
+            context,
+        };
+        compute_layout(&mut taffy_view, node, available_space)
     }
 }
 
diff --git a/src/util/debug.rs b/src/util/debug.rs
index 5c257533a..6ddffe7f9 100644
--- a/src/util/debug.rs
+++ b/src/util/debug.rs
@@ -1,20 +1,21 @@
 #![allow(dead_code)]
 
-use crate::tree::NodeId;
-use crate::{style, LayoutTree};
+use crate::prelude::Measurable;
+use crate::tree::{NodeId, Taffy};
+use crate::{style};
 use core::fmt::{Debug, Display, Write};
 use std::sync::Mutex;
 
 /// Prints a debug representation of the computed layout for a tree of nodes, starting with the passed root node.
-pub fn print_tree(tree: &impl LayoutTree, root: NodeId) {
+pub fn print_tree<M: Measurable>(tree: &Taffy<M>, root: NodeId) {
     println!("TREE");
     print_node(tree, root, false, String::new());
 }
 
-fn print_node(tree: &impl LayoutTree, node: NodeId, has_sibling: bool, lines_string: String) {
-    let layout = &tree.layout(node);
-    let style = &tree.style(node);
-    let num_children = tree.child_count(node);
+fn print_node<M: Measurable>(tree: &Taffy<M>, node: NodeId, has_sibling: bool, lines_string: String) {
+    let layout = &tree.layout(node).unwrap();
+    let style = &tree.style(node).unwrap();
+    let num_children = tree.child_count(node).unwrap();
 
     let display = match (num_children, style.display) {
         (_, style::Display::None) => "NONE",
@@ -43,7 +44,7 @@ fn print_node(tree: &impl LayoutTree, node: NodeId, has_sibling: bool, lines_str
     let new_string = lines_string + bar;
 
     // Recurse into children
-    for (index, child) in tree.children(node).enumerate() {
+    for (index, child) in tree.children(node).unwrap().into_iter().enumerate() {
         let has_sibling = index < num_children - 1;
         print_node(tree, child, has_sibling, new_string.clone());
     }