Fix child and child_count methods to use resolved children

DioxusLabs · Oct 12, 2023 · 1b65dac · 1b65dac
1 parent 6cb0783
commit 1b65dac
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 18 deletions.
diff --git a/src/compute/taffy_tree.rs b/src/compute/taffy_tree.rs
@@ -330,7 +330,7 @@ fn round_layout<NodeContext>(tree: &mut Taffy<NodeContext>, node_id: NodeId, cum
     layout.size.width = round(cumulative_x + unrounded_layout.size.width) - round(cumulative_x);
     layout.size.height = round(cumulative_y + unrounded_layout.size.height) - round(cumulative_y);
 
-    let child_count = tree.child_count(node_id).unwrap();
+    let child_count = tree.children[node_id.into()].len();
     for index in 0..child_count {
         let child = tree.children[node_id.into()][index];
         round_layout(tree, child, cumulative_x, cumulative_y);

diff --git a/src/tree/taffy_tree/tree.rs b/src/tree/taffy_tree/tree.rs
@@ -123,16 +123,10 @@ where
     pub(crate) fn new(taffy: &'t mut Taffy<NodeContext>, measure_function: MeasureFunction) -> Self {
         TaffyView { taffy, measure_function, node_children_cache: RefCell::new(ChildrenCache::new()) }
     }
-}
-
-impl<'t, NodeContext, MeasureFunction> LayoutTree for TaffyView<'t, NodeContext, MeasureFunction>
-where
-    MeasureFunction: FnMut(Size<Option<f32>>, Size<AvailableSpace>, NodeId, Option<&mut NodeContext>) -> Size<f32>,
-{
-    type ChildIter<'a> = RefCellVecIter<'a> where Self: 'a;
 
-    #[inline(always)]
-    fn children(&self, node: NodeId) -> Self::ChildIter<'_> {
+    /// Returns the resolved children, taking into account `Display::Contents`
+    /// Will use cached result if available, else compute and cache.
+    fn resolve_children(&self, node: NodeId) -> RefMut<'_, Vec<NodeId>> {
         let mut cache = self.node_children_cache.borrow_mut();
 
         // If the cache key does not match the requested node_id, then recompute the children for
@@ -144,12 +138,29 @@ where
         }
 
         // In all cases, return a reference into the cache
-        RefCellVecIter { children: RefMut::map(cache, |c| &mut c.children), index: 0 }
+        RefMut::map(cache, |c| &mut c.children)
+    }
+}
+
+impl<'t, NodeContext, MeasureFunction> LayoutTree for TaffyView<'t, NodeContext, MeasureFunction>
+where
+    MeasureFunction: FnMut(Size<Option<f32>>, Size<AvailableSpace>, NodeId, Option<&mut NodeContext>) -> Size<f32>,
+{
+    type ChildIter<'a> = RefCellVecIter<'a> where Self: 'a;
+
+    #[inline(always)]
+    fn children(&self, node: NodeId) -> Self::ChildIter<'_> {
+        RefCellVecIter { children: self.resolve_children(node), index: 0 }
     }
 
     #[inline(always)]
     fn child_count(&self, node: NodeId) -> usize {
-        self.taffy.children[node.into()].len()
+        self.resolve_children(node).len()
+    }
+
+    #[inline(always)]
+    fn child(&self, node: NodeId, id: usize) -> NodeId {
+        self.resolve_children(node)[id]
     }
 
     #[inline(always)]
@@ -175,12 +186,6 @@ where
         }
     }
 
-    #[inline(always)]
-    fn child(&self, node: NodeId, id: usize) -> NodeId {
-        self.taffy.children[node.into()][id]
-    }
-
-    #[inline(always)]
     fn measure_child_size(
         &mut self,
         node: NodeId,