From 783cccf95d3f9175a298eb604e97b378da1bedce Mon Sep 17 00:00:00 2001
From: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>
Date: Fri, 23 Aug 2024 18:59:30 +0200
Subject: [PATCH] WIP: Improve performance of Wrap Map (#16761)

We've ran into performance issues when reinserting new blocks into the
assistant panel; in profiles WrapMap showed up, as we try to query wrap
boundaries over and over, which is a hidden O(n^2) - for each block, we
may potentially look at all of the Wraps. This PR alleviates this issue
by storing away previously resolved wrap range; consecutive iterations
can often reuse it.
This should help with performance of Assistant Panel with long
conversations.

Release Notes:

- Improved performance of assistant panel with large # of text.
---
 crates/editor/src/display_map/block_map.rs | 41 ++++++++++++++++------
 1 file changed, 31 insertions(+), 10 deletions(-)

diff --git a/crates/editor/src/display_map/block_map.rs b/crates/editor/src/display_map/block_map.rs
index ba84ba48ac..020f2b5381 100644
--- a/crates/editor/src/display_map/block_map.rs
+++ b/crates/editor/src/display_map/block_map.rs
@@ -783,11 +783,13 @@ impl<'a> BlockMapWriter<'a> {
         &mut self,
         blocks: impl IntoIterator<Item = BlockProperties<Anchor>>,
     ) -> Vec<CustomBlockId> {
-        let mut ids = Vec::new();
+        let blocks = blocks.into_iter();
+        let mut ids = Vec::with_capacity(blocks.size_hint().1.unwrap_or(0));
         let mut edits = Patch::default();
         let wrap_snapshot = &*self.0.wrap_snapshot.borrow();
         let buffer = wrap_snapshot.buffer_snapshot();
 
+        let mut previous_wrap_row_range: Option<Range<u32>> = None;
         for block in blocks {
             let id = CustomBlockId(self.0.next_block_id.fetch_add(1, SeqCst));
             ids.push(id);
@@ -797,11 +799,18 @@ impl<'a> BlockMapWriter<'a> {
             let wrap_row = wrap_snapshot
                 .make_wrap_point(Point::new(point.row, 0), Bias::Left)
                 .row();
-            let start_row = wrap_snapshot.prev_row_boundary(WrapPoint::new(wrap_row, 0));
-            let end_row = wrap_snapshot
-                .next_row_boundary(WrapPoint::new(wrap_row, 0))
-                .unwrap_or(wrap_snapshot.max_point().row() + 1);
 
+            let (start_row, end_row) = {
+                previous_wrap_row_range.take_if(|range| !range.contains(&wrap_row));
+                let range = previous_wrap_row_range.get_or_insert_with(|| {
+                    let start_row = wrap_snapshot.prev_row_boundary(WrapPoint::new(wrap_row, 0));
+                    let end_row = wrap_snapshot
+                        .next_row_boundary(WrapPoint::new(wrap_row, 0))
+                        .unwrap_or(wrap_snapshot.max_point().row() + 1);
+                    start_row..end_row
+                });
+                (range.start, range.end)
+            };
             let block_ix = match self
                 .0
                 .custom_blocks
@@ -881,6 +890,7 @@ impl<'a> BlockMapWriter<'a> {
         let buffer = wrap_snapshot.buffer_snapshot();
         let mut edits = Patch::default();
         let mut last_block_buffer_row = None;
+        let mut previous_wrap_row_range: Option<Range<u32>> = None;
         self.0.custom_blocks.retain(|block| {
             if block_ids.contains(&block.id) {
                 let buffer_row = block.position.to_point(buffer).row;
@@ -889,21 +899,32 @@ impl<'a> BlockMapWriter<'a> {
                     let wrap_row = wrap_snapshot
                         .make_wrap_point(Point::new(buffer_row, 0), Bias::Left)
                         .row();
-                    let start_row = wrap_snapshot.prev_row_boundary(WrapPoint::new(wrap_row, 0));
-                    let end_row = wrap_snapshot
-                        .next_row_boundary(WrapPoint::new(wrap_row, 0))
-                        .unwrap_or(wrap_snapshot.max_point().row() + 1);
+                    let (start_row, end_row) = {
+                        previous_wrap_row_range.take_if(|range| !range.contains(&wrap_row));
+                        let range = previous_wrap_row_range.get_or_insert_with(|| {
+                            let start_row =
+                                wrap_snapshot.prev_row_boundary(WrapPoint::new(wrap_row, 0));
+                            let end_row = wrap_snapshot
+                                .next_row_boundary(WrapPoint::new(wrap_row, 0))
+                                .unwrap_or(wrap_snapshot.max_point().row() + 1);
+                            start_row..end_row
+                        });
+                        (range.start, range.end)
+                    };
+
                     edits.push(Edit {
                         old: start_row..end_row,
                         new: start_row..end_row,
                     })
                 }
-                self.0.custom_blocks_by_id.remove(&block.id);
                 false
             } else {
                 true
             }
         });
+        self.0
+            .custom_blocks_by_id
+            .retain(|id, _| !block_ids.contains(id));
         self.0.sync(wrap_snapshot, edits);
     }
 }