From bafd7ed000fd8b370fc354d482ed82f0aef20ace Mon Sep 17 00:00:00 2001
From: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>
Date: Fri, 4 Oct 2024 13:58:57 +0200
Subject: [PATCH] gpui: Store measure functions as context of taffy nodes
 (#18732)

Taffy maintains a mapping of NodeId <-> Context anyways (and does the
lookup), so it's redundant for us to store it separately. Tl;dr: we get
rid of one map and one map lookup per layout request.

Release Notes:

- N/A
---
 crates/gpui/src/taffy.rs | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/crates/gpui/src/taffy.rs b/crates/gpui/src/taffy.rs
index feec5bf8c1..9d99078ae5 100644
--- a/crates/gpui/src/taffy.rs
+++ b/crates/gpui/src/taffy.rs
@@ -15,12 +15,13 @@ use taffy::{
 type NodeMeasureFn =
     Box<dyn FnMut(Size<Option<Pixels>>, Size<AvailableSpace>, &mut WindowContext) -> Size<Pixels>>;
 
+struct NodeContext {
+    measure: NodeMeasureFn,
+}
 pub struct TaffyLayoutEngine {
-    taffy: TaffyTree<()>,
-    styles: FxHashMap<LayoutId, Style>,
+    taffy: TaffyTree<NodeContext>,
     absolute_layout_bounds: FxHashMap<LayoutId, Bounds<Pixels>>,
     computed_layouts: FxHashSet<LayoutId>,
-    nodes_to_measure: FxHashMap<LayoutId, NodeMeasureFn>,
 }
 
 const EXPECT_MESSAGE: &str = "we should avoid taffy layout errors by construction if possible";
@@ -29,10 +30,8 @@ impl TaffyLayoutEngine {
     pub fn new() -> Self {
         TaffyLayoutEngine {
             taffy: TaffyTree::new(),
-            styles: FxHashMap::default(),
             absolute_layout_bounds: FxHashMap::default(),
             computed_layouts: FxHashSet::default(),
-            nodes_to_measure: FxHashMap::default(),
         }
     }
 
@@ -40,8 +39,6 @@ impl TaffyLayoutEngine {
         self.taffy.clear();
         self.absolute_layout_bounds.clear();
         self.computed_layouts.clear();
-        self.nodes_to_measure.clear();
-        self.styles.clear();
     }
 
     pub fn request_layout(
@@ -67,7 +64,6 @@ impl TaffyLayoutEngine {
                 .into();
             parent_id
         };
-        self.styles.insert(layout_id, style);
         layout_id
     }
 
@@ -82,11 +78,14 @@ impl TaffyLayoutEngine {
 
         let layout_id = self
             .taffy
-            .new_leaf_with_context(taffy_style, ())
+            .new_leaf_with_context(
+                taffy_style,
+                NodeContext {
+                    measure: Box::new(measure),
+                },
+            )
             .expect(EXPECT_MESSAGE)
             .into();
-        self.nodes_to_measure.insert(layout_id, Box::new(measure));
-        self.styles.insert(layout_id, style);
         layout_id
     }
 
@@ -175,8 +174,8 @@ impl TaffyLayoutEngine {
             .compute_layout_with_measure(
                 id.into(),
                 available_space.into(),
-                |known_dimensions, available_space, node_id, _context, _style| {
-                    let Some(measure) = self.nodes_to_measure.get_mut(&node_id.into()) else {
+                |known_dimensions, available_space, _node_id, node_context, _style| {
+                    let Some(node_context) = node_context else {
                         return taffy::geometry::Size::default();
                     };
 
@@ -185,7 +184,7 @@ impl TaffyLayoutEngine {
                         height: known_dimensions.height.map(Pixels),
                     };
 
-                    measure(known_dimensions, available_space.into(), cx).into()
+                    (node_context.measure)(known_dimensions, available_space.into(), cx).into()
                 },
             )
             .expect(EXPECT_MESSAGE);