1 files changed, 61 insertions, 34 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] {