From d9b8c07b1e2bac2185dd5b4a08e210a7fc2b9e56 Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Thu, 30 Nov 2023 16:03:22 +0100
Subject: [PATCH] Lift Send and Sync restriction for measurements

---
 crates/gpui2/src/elements/text.rs         |  4 +--
 crates/gpui2/src/elements/uniform_list.rs |  2 --
 crates/gpui2/src/taffy.rs                 | 33 ++++++++++++++---------
 crates/gpui2/src/window.rs                |  4 +--
 4 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/crates/gpui2/src/elements/text.rs b/crates/gpui2/src/elements/text.rs
index d57daca062..2f432bcb75 100644
--- a/crates/gpui2/src/elements/text.rs
+++ b/crates/gpui2/src/elements/text.rs
@@ -152,15 +152,13 @@ impl TextState {
             .to_pixels(font_size.into(), cx.rem_size());
         let text = SharedString::from(text);
 
-        let rem_size = cx.rem_size();
-
         let runs = if let Some(runs) = runs {
             runs
         } else {
             vec![text_style.to_run(text.len())]
         };
 
-        let layout_id = cx.request_measured_layout(Default::default(), rem_size, {
+        let layout_id = cx.request_measured_layout(Default::default(), {
             let element_state = self.clone();
 
             move |known_dimensions, available_space| {
diff --git a/crates/gpui2/src/elements/uniform_list.rs b/crates/gpui2/src/elements/uniform_list.rs
index 2d5a46f3d9..ed29caa43d 100644
--- a/crates/gpui2/src/elements/uniform_list.rs
+++ b/crates/gpui2/src/elements/uniform_list.rs
@@ -109,7 +109,6 @@ impl Element for UniformList {
         cx: &mut WindowContext,
     ) -> (LayoutId, Self::State) {
         let max_items = self.item_count;
-        let rem_size = cx.rem_size();
         let item_size = state
             .as_ref()
             .map(|s| s.item_size)
@@ -120,7 +119,6 @@ impl Element for UniformList {
                 .layout(state.map(|s| s.interactive), cx, |style, cx| {
                     cx.request_measured_layout(
                         style,
-                        rem_size,
                         move |known_dimensions: Size<Option<Pixels>>,
                               available_space: Size<AvailableSpace>| {
                             let desired_height = item_size.height * max_items;
diff --git a/crates/gpui2/src/taffy.rs b/crates/gpui2/src/taffy.rs
index 81a057055a..7e25b5c5b2 100644
--- a/crates/gpui2/src/taffy.rs
+++ b/crates/gpui2/src/taffy.rs
@@ -9,13 +9,15 @@ use taffy::{
     Taffy,
 };
 
-type Measureable = dyn Fn(Size<Option<Pixels>>, Size<AvailableSpace>) -> Size<Pixels> + Send + Sync;
-
 pub struct TaffyLayoutEngine {
-    taffy: Taffy<Box<Measureable>>,
+    taffy: Taffy,
     children_to_parents: HashMap<LayoutId, LayoutId>,
     absolute_layout_bounds: HashMap<LayoutId, Bounds<Pixels>>,
     computed_layouts: HashSet<LayoutId>,
+    nodes_to_measure: HashMap<
+        LayoutId,
+        Box<dyn FnMut(Size<Option<Pixels>>, Size<AvailableSpace>) -> Size<Pixels>>,
+    >,
 }
 
 static EXPECT_MESSAGE: &'static str =
@@ -28,6 +30,7 @@ impl TaffyLayoutEngine {
             children_to_parents: HashMap::default(),
             absolute_layout_bounds: HashMap::default(),
             computed_layouts: HashSet::default(),
+            nodes_to_measure: HashMap::default(),
         }
     }
 
@@ -36,6 +39,7 @@ impl TaffyLayoutEngine {
         self.children_to_parents.clear();
         self.absolute_layout_bounds.clear();
         self.computed_layouts.clear();
+        self.nodes_to_measure.clear();
     }
 
     pub fn request_layout(
@@ -65,18 +69,17 @@ impl TaffyLayoutEngine {
         &mut self,
         style: Style,
         rem_size: Pixels,
-        measure: impl Fn(Size<Option<Pixels>>, Size<AvailableSpace>) -> Size<Pixels>
-            + Send
-            + Sync
-            + 'static,
+        measure: impl FnMut(Size<Option<Pixels>>, Size<AvailableSpace>) -> Size<Pixels> + 'static,
     ) -> LayoutId {
         let style = style.to_taffy(rem_size);
 
-        let measurable = Box::new(measure);
-        self.taffy
-            .new_leaf_with_context(style, measurable)
+        let layout_id = self
+            .taffy
+            .new_leaf_with_context(style, ())
             .expect(EXPECT_MESSAGE)
-            .into()
+            .into();
+        self.nodes_to_measure.insert(layout_id, Box::new(measure));
+        layout_id
     }
 
     // Used to understand performance
@@ -155,12 +158,14 @@ impl TaffyLayoutEngine {
         }
 
         // let started_at = std::time::Instant::now();
+        dbg!(">>>>>>>>>>>>>");
+
         self.taffy
             .compute_layout_with_measure(
                 id.into(),
                 available_space.into(),
-                |known_dimensions, available_space, _node_id, context| {
-                    let Some(measure) = context else {
+                |known_dimensions, available_space, node_id, _context| {
+                    let Some(measure) = self.nodes_to_measure.get_mut(&node_id.into()) else {
                         return taffy::geometry::Size::default();
                     };
 
@@ -173,6 +178,8 @@ impl TaffyLayoutEngine {
                 },
             )
             .expect(EXPECT_MESSAGE);
+
+        dbg!("<<<<<<");
         // println!("compute_layout took {:?}", started_at.elapsed());
     }
 
diff --git a/crates/gpui2/src/window.rs b/crates/gpui2/src/window.rs
index 76932f28e4..6f0d1381d0 100644
--- a/crates/gpui2/src/window.rs
+++ b/crates/gpui2/src/window.rs
@@ -618,13 +618,13 @@ impl<'a> WindowContext<'a> {
     /// The given closure is invoked at layout time with the known dimensions and available space and
     /// returns a `Size`.
     pub fn request_measured_layout<
-        F: Fn(Size<Option<Pixels>>, Size<AvailableSpace>) -> Size<Pixels> + Send + Sync + 'static,
+        F: FnMut(Size<Option<Pixels>>, Size<AvailableSpace>) -> Size<Pixels> + 'static,
     >(
         &mut self,
         style: Style,
-        rem_size: Pixels,
         measure: F,
     ) -> LayoutId {
+        let rem_size = self.rem_size();
         self.window
             .layout_engine
             .request_measured_layout(style, rem_size, measure)