map positions through changes in O(N)

4 files changed, 129 insertions, 72 deletions

diff --git a/helix-core/src/selection.rs b/helix-core/src/selection.rs
index 9e366c33..0539c9f6 100644
--- a/helix-core/src/selection.rs
+++ b/helix-core/src/selection.rs
@@ -161,34 +161,35 @@ impl Range {
         self.from() <= pos && pos < self.to()
     }
 
-    /// Map a range through a set of changes. Returns a new range representing the same position
-    /// after the changes are applied.
-    pub fn map(self, changes: &ChangeSet) -> Self {
+    /// Map a range through a set of changes. Returns a new range representing
+    /// the same position after the changes are applied. Note that this
+    /// function runs in O(N) (N is number of changes) and can therefore
+    /// cause performance problems if run for a large number of ranges as the
+    /// complexity is then O(MN) (for multicuror M=N usually). Instead use
+    /// [Selection::map] or [ChangeSet::update_positions] instead
+    pub fn map(mut self, changes: &ChangeSet) -> Self {
         use std::cmp::Ordering;
-        let (anchor, head) = match self.anchor.cmp(&self.head) {
-            Ordering::Equal => (
-                changes.map_pos(self.anchor, Assoc::After),
-                changes.map_pos(self.head, Assoc::After),
-            ),
-            Ordering::Less => (
-                changes.map_pos(self.anchor, Assoc::After),
-                changes.map_pos(self.head, Assoc::Before),
-            ),
-            Ordering::Greater => (
-                changes.map_pos(self.anchor, Assoc::Before),
-                changes.map_pos(self.head, Assoc::After),
-            ),
-        };
-
-        // We want to return a new `Range` with `horiz == None` every time,
-        // even if the anchor and head haven't changed, because we don't
-        // know if the *visual* position hasn't changed due to
-        // character-width or grapheme changes earlier in the text.
-        Self {
-            anchor,
-            head,
-            old_visual_position: None,
+        if changes.is_empty() {
+            return self;
         }
+
+        let positions_to_map = match self.anchor.cmp(&self.head) {
+            Ordering::Equal => [
+                (&mut self.anchor, Assoc::After),
+                (&mut self.head, Assoc::After),
+            ],
+            Ordering::Less => [
+                (&mut self.anchor, Assoc::After),
+                (&mut self.head, Assoc::Before),
+            ],
+            Ordering::Greater => [
+                (&mut self.head, Assoc::After),
+                (&mut self.anchor, Assoc::Before),
+            ],
+        };
+        changes.update_positions(positions_to_map.into_iter());
+        self.old_visual_position = None;
+        self
     }
 
     /// Extend the range to cover at least `from` `to`.
@@ -450,18 +451,44 @@ impl Selection {
 
     /// Map selections over a set of changes. Useful for adjusting the selection position after
     /// applying changes to a document.
-    pub fn map(self, changes: &ChangeSet) -> Self {
+    pub fn map(mut self, changes: &ChangeSet) -> Self {
         if changes.is_empty() {
             return self;
         }
+        self = self.map_no_normalize(changes);
+        // TODO: only normalize if needed (any ranges out of order)
+        self = self.normalize();
 
-        Self::new(
-            self.ranges
-                .into_iter()
-                .map(|range| range.map(changes))
-                .collect(),
-            self.primary_index,
-        )
+        self
+    }
+
+    /// Map selections over a set of changes. Useful for adjusting the selection position after
+    /// applying changes to a document. Doesn't normalize the selection
+    pub fn map_no_normalize(mut self, changes: &ChangeSet) -> Self {
+        if changes.is_empty() {
+            return self;
+        }
+
+        let positions_to_map = self.ranges.iter_mut().flat_map(|range| {
+            use std::cmp::Ordering;
+            range.old_visual_position = None;
+            match range.anchor.cmp(&range.head) {
+                Ordering::Equal => [
+                    (&mut range.anchor, Assoc::After),
+                    (&mut range.head, Assoc::After),
+                ],
+                Ordering::Less => [
+                    (&mut range.anchor, Assoc::After),
+                    (&mut range.head, Assoc::Before),
+                ],
+                Ordering::Greater => [
+                    (&mut range.head, Assoc::After),
+                    (&mut range.anchor, Assoc::Before),
+                ],
+            }
+        });
+        changes.update_positions(positions_to_map);
+        self
     }
 
     pub fn ranges(&self) -> &[Range] {
diff --git a/helix-core/src/transaction.rs b/helix-core/src/transaction.rs
index f4f94b54..53e53c23 100644
--- a/helix-core/src/transaction.rs
+++ b/helix-core/src/transaction.rs
@@ -1,7 +1,7 @@
 use smallvec::SmallVec;
 
 use crate::{Range, Rope, Selection, Tendril};
-use std::borrow::Cow;
+use std::{borrow::Cow, iter::once};
 
 /// (from, to, replacement)
 pub type Change = (usize, usize, Option<Tendril>);
@@ -326,17 +326,33 @@ impl ChangeSet {
         self.changes.is_empty() || self.changes == [Operation::Retain(self.len)]
     }
 
-    /// Map a position through the changes.
+    /// Map a *sorted* list of positions through the changes.
     ///
-    /// `assoc` indicates which size to associate the position with. `Before` will keep the
-    /// position close to the character before, and will place it before insertions over that
-    /// range, or at that point. `After` will move it forward, placing it at the end of such
-    /// insertions.
-    pub fn map_pos(&self, pos: usize, assoc: Assoc) -> usize {
+    /// This is equivalent to updating each position with `map_pos`:
+    ///
+    /// ``` no-compile
+    /// for (pos, assoc) in positions {
+    ///     *pos = changes.map_pos(*pos, assoc);
+    /// }
+    /// ```
+    /// However this function is significantly faster running in `O(N+M)` instead of `O(NM)`
+    pub fn update_positions<'a>(&self, positions: impl Iterator<Item = (&'a mut usize, Assoc)>) {
         use Operation::*;
+
+        let mut positions = positions.peekable();
+        macro_rules! map {
+            ($map: expr) => {
+                loop {
+                    let Some((pos, assoc)) = positions.peek_mut() else { return; };
+                    let Some(new_pos) = $map(**pos, *assoc) else { break; };
+                    **pos = new_pos;
+                    positions.next();
+                }
+            };
+        }
+
         let mut old_pos = 0;
         let mut new_pos = 0;
-
         let mut iter = self.changes.iter().peekable();
 
         while let Some(change) = iter.next() {
@@ -348,16 +364,12 @@ impl ChangeSet {
 
             match change {
                 Retain(_) => {
-                    if old_end > pos {
-                        return new_pos + (pos - old_pos);
-                    }
+                    map!(|pos, _| (old_end > pos).then_some(new_pos + (pos - old_pos)));
                     new_pos += len;
                 }
                 Delete(_) => {
                     // in range
-                    if old_end > pos {
-                        return new_pos;
-                    }
+                    map!(|pos, _| (old_end > pos).then_some(new_pos));
                 }
                 Insert(s) => {
                     let ins = s.chars().count();
@@ -368,26 +380,26 @@ impl ChangeSet {
 
                         old_end = old_pos + len;
                         // in range of replaced text
-                        if old_end > pos {
+                        map!(|pos, assoc| (old_end > pos).then(|| {
                             // at point or tracking before
                             if pos == old_pos || assoc == Assoc::Before {
-                                return new_pos;
+                                new_pos
                             } else {
                                 // place to end of insert
-                                return new_pos + ins;
+                                new_pos + ins
                             }
-                        }
+                        }));
                     } else {
                         // at insert point
-                        if old_pos == pos {
+                        map!(|pos, assoc| (old_pos == pos).then(|| {
                             // return position before inserted text
                             if assoc == Assoc::Before {
-                                return new_pos;
+                                new_pos
                             } else {
                                 // after text
-                                return new_pos + ins;
+                                new_pos + ins
                             }
-                        }
+                        }));
                     }
 
                     new_pos += ins;
@@ -395,14 +407,21 @@ impl ChangeSet {
             }
             old_pos = old_end;
         }
+        map!(|pos, _| (old_pos == pos).then_some(new_pos));
+        let out_of_bounds: Vec<_> = positions.collect();
 
-        if pos > old_pos {
-            panic!(
-                "Position {} is out of range for changeset len {}!",
-                pos, old_pos
-            )
-        }
-        new_pos
+        panic!("Positions {out_of_bounds:?} are out of range for changeset len {old_pos}!",)
+    }
+
+    /// Map a position through the changes.
+    ///
+    /// `assoc` indicates which side to associate the position with. `Before` will keep the
+    /// position close to the character before, and will place it before insertions over that
+    /// range, or at that point. `After` will move it forward, placing it at the end of such
+    /// insertions.
+    pub fn map_pos(&self, mut pos: usize, assoc: Assoc) -> usize {
+        self.update_positions(once((&mut pos, assoc)));
+        pos
     }
 
     pub fn changes_iter(&self) -> ChangeIterator {
diff --git a/helix-lsp/src/lib.rs b/helix-lsp/src/lib.rs
index d053dbf9..277a4c28 100644
--- a/helix-lsp/src/lib.rs
+++ b/helix-lsp/src/lib.rs
@@ -377,7 +377,7 @@ pub mod util {
         .expect("transaction must be valid for primary selection");
         let removed_text = text.slice(removed_start..removed_end);
 
-        let (transaction, selection) = Transaction::change_by_selection_ignore_overlapping(
+        let (transaction, mut selection) = Transaction::change_by_selection_ignore_overlapping(
             doc,
             selection,
             |range| {
@@ -420,6 +420,11 @@ pub mod util {
             return transaction;
         }
 
+        // Don't normalize to avoid merging/reording selections which would
+        // break the association between tabstops and selections. Most ranges
+        // will be replaced by tabstops anyways and the final selection will be
+        // normalized anyways
+        selection = selection.map_no_normalize(changes);
         let mut mapped_selection = SmallVec::with_capacity(selection.len());
         let mut mapped_primary_idx = 0;
         let primary_range = selection.primary();
@@ -428,7 +433,6 @@ pub mod util {
                 mapped_primary_idx = mapped_selection.len()
             }
 
-            let range = range.map(changes);
             let tabstops = tabstops.first().filter(|tabstops| !tabstops.is_empty());
             let Some(tabstops) = tabstops else{
                 // no tabstop normal mapping
diff --git a/helix-view/src/document.rs b/helix-view/src/document.rs
index 52b2dfc0..4140c105 100644
--- a/helix-view/src/document.rs
+++ b/helix-view/src/document.rs
@@ -1121,21 +1121,28 @@ impl Document {
 
             let changes = transaction.changes();
 
+            changes.update_positions(
+                self.diagnostics
+                    .iter_mut()
+                    .map(|diagnostic| (&mut diagnostic.range.start, Assoc::After)),
+            );
+            changes.update_positions(
+                self.diagnostics
+                    .iter_mut()
+                    .map(|diagnostic| (&mut diagnostic.range.end, Assoc::After)),
+            );
             // map state.diagnostics over changes::map_pos too
             for diagnostic in &mut self.diagnostics {
-                diagnostic.range.start = changes.map_pos(diagnostic.range.start, Assoc::After);
-                diagnostic.range.end = changes.map_pos(diagnostic.range.end, Assoc::After);
                 diagnostic.line = self.text.char_to_line(diagnostic.range.start);
             }
-            self.diagnostics
-                .sort_unstable_by_key(|diagnostic| diagnostic.range);
 
             // Update the inlay hint annotations' positions, helping ensure they are displayed in the proper place
             let apply_inlay_hint_changes = |annotations: &mut Rc<[InlineAnnotation]>| {
                 if let Some(data) = Rc::get_mut(annotations) {
-                    for inline in data.iter_mut() {
-                        inline.char_idx = changes.map_pos(inline.char_idx, Assoc::After);
-                    }
+                    changes.update_positions(
+                        data.iter_mut()
+                            .map(|diagnostic| (&mut diagnostic.char_idx, Assoc::After)),
+                    );
                 }
             };