From 0ae522f3df433bb778fa2ff98fd3d7047021c6ef Mon Sep 17 00:00:00 2001
From: Nathan Vegdahl <cessen@cessen.com>
Date: Mon, 28 Jun 2021 11:40:07 -0700
Subject: [PATCH] Clean up `Selection` to not use so many allocations.

---
 helix-core/src/object.rs    |   2 +-
 helix-core/src/selection.rs | 154 ++++++++++++++++-------------
 helix-term/src/commands.rs  | 187 ++++++++++++------------------------
 helix-view/src/document.rs  |  20 +++-
 4 files changed, 168 insertions(+), 195 deletions(-)

diff --git a/helix-core/src/object.rs b/helix-core/src/object.rs
index 1c644fb2..863b6e55 100644
--- a/helix-core/src/object.rs
+++ b/helix-core/src/object.rs
@@ -6,7 +6,7 @@ use smallvec::smallvec;
 pub fn expand_selection(syntax: &Syntax, text: RopeSlice, selection: &Selection) -> Selection {
     let tree = syntax.tree();
 
-    selection.transform(|range| {
+    selection.clone().transformed(|range| {
         let from = text.char_to_byte(range.from());
         let to = text.char_to_byte(range.to());
 
diff --git a/helix-core/src/selection.rs b/helix-core/src/selection.rs
index 4260c002..ebc88f8b 100644
--- a/helix-core/src/selection.rs
+++ b/helix-core/src/selection.rs
@@ -137,6 +137,27 @@ impl Range {
         }
     }
 
+    /// Returns a range that encompasses both input ranges.
+    ///
+    /// This is like `extend()`, but tries to negotiate the
+    /// anchor/head ordering between the two input ranges.
+    #[must_use]
+    pub fn merge(&self, other: Self) -> Self {
+        if self.anchor > self.head && other.anchor > other.head {
+            Range {
+                anchor: self.anchor.max(other.anchor),
+                head: self.head.min(other.head),
+                horiz: None,
+            }
+        } else {
+            Range {
+                anchor: self.from().min(other.from()),
+                head: self.to().max(other.to()),
+                horiz: None,
+            }
+        }
+    }
+
     /// Compute a possibly new range from this range, attempting to ensure
     /// a minimum range width of 1 char by shifting the head in the forward
     /// direction as needed.
@@ -170,19 +191,20 @@ impl Range {
     /// ranges will never collapse to zero-width.
     #[must_use]
     pub fn grapheme_aligned(&self, slice: RopeSlice) -> Self {
-        let (new_anchor, new_head) = if self.anchor == self.head {
-            let pos = ensure_grapheme_boundary_prev(slice, self.anchor);
-            (pos, pos)
-        } else if self.anchor < self.head {
-            (
+        use std::cmp::Ordering;
+        let (new_anchor, new_head) = match self.anchor.cmp(&self.head) {
+            Ordering::Equal => {
+                let pos = ensure_grapheme_boundary_prev(slice, self.anchor);
+                (pos, pos)
+            }
+            Ordering::Less => (
                 ensure_grapheme_boundary_prev(slice, self.anchor),
                 ensure_grapheme_boundary_next(slice, self.head),
-            )
-        } else {
-            (
+            ),
+            Ordering::Greater => (
                 ensure_grapheme_boundary_next(slice, self.anchor),
                 ensure_grapheme_boundary_prev(slice, self.head),
-            )
+            ),
         };
         Range {
             anchor: new_anchor,
@@ -238,10 +260,8 @@ impl Selection {
     }
 
     pub fn push(mut self, range: Range) -> Self {
-        let index = self.ranges.len();
         self.ranges.push(range);
-
-        Self::normalize(self.ranges, index)
+        self.normalized()
     }
     // replace_range
 
@@ -287,80 +307,80 @@ impl Selection {
         Self::single(pos, pos)
     }
 
-    fn normalize(mut ranges: SmallVec<[Range; 1]>, mut primary_index: usize) -> Self {
-        let primary = ranges[primary_index];
-        ranges.sort_unstable_by_key(Range::from);
-        primary_index = ranges.iter().position(|&range| range == primary).unwrap();
+    /// Normalizes a `Selection`.
+    pub fn normalize(&mut self) -> &mut Self {
+        let primary = self.ranges[self.primary_index];
+        self.ranges.sort_unstable_by_key(Range::from);
+        self.primary_index = self
+            .ranges
+            .iter()
+            .position(|&range| range == primary)
+            .unwrap();
 
-        let mut result = SmallVec::with_capacity(ranges.len()); // approx
-
-        // TODO: we could do with one vec by removing elements as we mutate
-
-        let mut i = 0;
-
-        for range in ranges.into_iter() {
-            // if previous value exists
-            if let Some(prev) = result.last_mut() {
-                // and we overlap it
-
-                // TODO: we used to simply check range.from() <(=) prev.to()
-                // avoiding two comparisons
-                if range.overlaps(prev) {
-                    let from = prev.from();
-                    let to = std::cmp::max(range.to(), prev.to());
-
-                    if i <= primary_index {
-                        primary_index -= 1
-                    }
-
-                    // merge into previous
-                    if range.anchor > range.head {
-                        prev.anchor = to;
-                        prev.head = from;
-                    } else {
-                        prev.anchor = from;
-                        prev.head = to;
-                    }
-                    continue;
+        let mut prev_i = 0;
+        for i in 1..self.ranges.len() {
+            if self.ranges[prev_i].overlaps(&self.ranges[i]) {
+                if i == self.primary_index {
+                    self.primary_index = prev_i;
                 }
+                self.ranges[prev_i] = self.ranges[prev_i].merge(self.ranges[i]);
+            } else {
+                prev_i += 1;
+                self.ranges[prev_i] = self.ranges[i];
             }
-
-            result.push(range);
-            i += 1
         }
 
-        Self {
-            ranges: result,
-            primary_index,
-        }
+        self.ranges.truncate(prev_i + 1);
+
+        self
+    }
+
+    /// Normalizes a `Selection`.
+    #[must_use]
+    pub fn normalized(mut self) -> Self {
+        self.normalize();
+        self
     }
 
     // TODO: consume an iterator or a vec to reduce allocations?
     #[must_use]
     pub fn new(ranges: SmallVec<[Range; 1]>, primary_index: usize) -> Self {
         assert!(!ranges.is_empty());
+        debug_assert!(primary_index < ranges.len());
 
-        // fast path for a single selection (cursor)
-        if ranges.len() == 1 {
-            return Self {
-                ranges,
-                primary_index: 0,
-            };
+        let mut selection = Self {
+            ranges,
+            primary_index,
+        };
+
+        if selection.ranges.len() > 1 {
+            // TODO: only normalize if needed (any ranges out of order)
+            selection.normalize();
         }
 
-        // TODO: only normalize if needed (any ranges out of order)
-        Self::normalize(ranges, primary_index)
+        selection
     }
 
-    /// Takes a closure and maps each selection over the closure.
-    pub fn transform<F>(&self, f: F) -> Self
+    /// Takes a closure and maps each `Range` over the closure.
+    pub fn transform<F>(&mut self, f: F) -> &mut Self
     where
         F: Fn(Range) -> Range,
     {
-        Self::new(
-            self.ranges.iter().copied().map(f).collect(),
-            self.primary_index,
-        )
+        for range in self.ranges.iter_mut() {
+            *range = f(*range)
+        }
+
+        self
+    }
+
+    /// Takes a closure and maps each `Range` over the closure.
+    #[must_use]
+    pub fn transformed<F>(mut self, f: F) -> Self
+    where
+        F: Fn(Range) -> Range,
+    {
+        self.transform(f);
+        self
     }
 
     pub fn fragments<'a>(&'a self, text: RopeSlice<'a>) -> impl Iterator<Item = Cow<str>> + 'a {
diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index fa251ff0..29a10c76 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -310,48 +310,44 @@ impl PartialEq for Command {
 fn move_char_left(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
-    let selection = doc.selection(view.id).transform(|range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         movement::move_horizontally(text, range, Direction::Backward, count, Movement::Move)
     });
-    doc.set_selection(view.id, selection);
 }
 
 fn move_char_right(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
-    let selection = doc.selection(view.id).transform(|range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         movement::move_horizontally(text, range, Direction::Forward, count, Movement::Move)
     });
-    doc.set_selection(view.id, selection);
 }
 
 fn move_line_up(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
-    let selection = doc.selection(view.id).transform(|range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         movement::move_vertically(text, range, Direction::Backward, count, Movement::Move)
     });
-    doc.set_selection(view.id, selection);
 }
 
 fn move_line_down(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
-    let selection = doc.selection(view.id).transform(|range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         movement::move_vertically(text, range, Direction::Forward, count, Movement::Move)
     });
-    doc.set_selection(view.id, selection);
 }
 
 fn move_line_end(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
 
-    let selection = doc.selection(view.id).transform(|range| {
-        let text = doc.text();
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         let line = text.char_to_line(range.head);
 
         let pos = line_end_char_index(&text.slice(..), line);
@@ -360,30 +356,26 @@ fn move_line_end(cx: &mut Context) {
 
         Range::new(pos, pos)
     });
-
-    doc.set_selection(view.id, selection);
 }
 
 fn move_line_start(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
 
-    let selection = doc.selection(view.id).transform(|range| {
-        let text = doc.text();
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         let line = text.char_to_line(range.head);
 
         // adjust to start of the line
         let pos = text.line_to_char(line);
         Range::new(pos, pos)
     });
-
-    doc.set_selection(view.id, selection);
 }
 
 fn move_first_nonwhitespace(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
 
-    let selection = doc.selection(view.id).transform(|range| {
-        let text = doc.text();
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         let line_idx = text.char_to_line(range.head);
 
         if let Some(pos) = find_first_non_whitespace_char(text.line(line_idx)) {
@@ -393,8 +385,6 @@ fn move_first_nonwhitespace(cx: &mut Context) {
             range
         }
     });
-
-    doc.set_selection(view.id, selection);
 }
 
 // TODO: move vs extend could take an extra type Extend/Move that would
@@ -404,73 +394,49 @@ fn move_first_nonwhitespace(cx: &mut Context) {
 fn move_next_word_start(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
 
-    let selection = doc
-        .selection(view.id)
-        .transform(|range| movement::move_next_word_start(text, range, count));
-
-    doc.set_selection(view.id, selection);
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| movement::move_next_word_start(text, range, count));
 }
 
 fn move_prev_word_start(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
 
-    let selection = doc
-        .selection(view.id)
-        .transform(|range| movement::move_prev_word_start(text, range, count));
-
-    doc.set_selection(view.id, selection);
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| movement::move_prev_word_start(text, range, count));
 }
 
 fn move_next_word_end(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
 
-    let selection = doc
-        .selection(view.id)
-        .transform(|range| movement::move_next_word_end(text, range, count));
-
-    doc.set_selection(view.id, selection);
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| movement::move_next_word_end(text, range, count));
 }
 
 fn move_next_long_word_start(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
 
-    let selection = doc
-        .selection(view.id)
-        .transform(|range| movement::move_next_long_word_start(text, range, count));
-
-    doc.set_selection(view.id, selection);
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| movement::move_next_long_word_start(text, range, count));
 }
 
 fn move_prev_long_word_start(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
 
-    let selection = doc
-        .selection(view.id)
-        .transform(|range| movement::move_prev_long_word_start(text, range, count));
-
-    doc.set_selection(view.id, selection);
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| movement::move_prev_long_word_start(text, range, count));
 }
 
 fn move_next_long_word_end(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
 
-    let selection = doc
-        .selection(view.id)
-        .transform(|range| movement::move_next_long_word_end(text, range, count));
-
-    doc.set_selection(view.id, selection);
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| movement::move_next_long_word_end(text, range, count));
 }
 
 fn move_file_start(cx: &mut Context) {
@@ -490,42 +456,37 @@ fn move_file_end(cx: &mut Context) {
 fn extend_next_word_start(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
 
-    let selection = doc.selection(view.id).transform(|mut range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|mut range| {
         let word = movement::move_next_word_start(text, range, count);
         let pos = word.head;
         Range::new(range.anchor, pos)
     });
-
-    doc.set_selection(view.id, selection);
 }
 
 fn extend_prev_word_start(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
 
-    let selection = doc.selection(view.id).transform(|mut range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|mut range| {
         let word = movement::move_prev_word_start(text, range, count);
         let pos = word.head;
         Range::new(range.anchor, pos)
     });
-    doc.set_selection(view.id, selection);
 }
 
 fn extend_next_word_end(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
 
-    let selection = doc.selection(view.id).transform(|mut range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|mut range| {
         let word = movement::move_next_word_end(text, range, count);
         let pos = word.head;
         Range::new(range.anchor, pos)
     });
-
-    doc.set_selection(view.id, selection);
 }
 
 #[inline]
@@ -568,9 +529,9 @@ where
         };
 
         let (view, doc) = current!(cx.editor);
-        let text = doc.text().slice(..);
 
-        let selection = doc.selection(view.id).transform(|mut range| {
+        let (text, selection) = doc.text_and_mut_selection(view.id);
+        selection.transform(|mut range| {
             search_fn(text, ch, range.head, count, inclusive).map_or(range, |pos| {
                 if extend {
                     Range::new(range.anchor, pos)
@@ -581,8 +542,6 @@ where
                 // or (pos, pos) to move to found val
             })
         });
-
-        doc.set_selection(view.id, selection);
     })
 }
 
@@ -661,8 +620,8 @@ fn extend_prev_char(cx: &mut Context) {
 fn extend_first_nonwhitespace(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
 
-    let selection = doc.selection(view.id).transform(|range| {
-        let text = doc.text();
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         let line_idx = text.char_to_line(range.head);
 
         if let Some(pos) = find_first_non_whitespace_char(text.line(line_idx)) {
@@ -672,8 +631,6 @@ fn extend_first_nonwhitespace(cx: &mut Context) {
             range
         }
     });
-
-    doc.set_selection(view.id, selection);
 }
 
 fn replace(cx: &mut Context) {
@@ -784,48 +741,44 @@ fn half_page_down(cx: &mut Context) {
 fn extend_char_left(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
-    let selection = doc.selection(view.id).transform(|range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         movement::move_horizontally(text, range, Direction::Backward, count, Movement::Extend)
     });
-    doc.set_selection(view.id, selection);
 }
 
 fn extend_char_right(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
-    let selection = doc.selection(view.id).transform(|range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         movement::move_horizontally(text, range, Direction::Forward, count, Movement::Extend)
     });
-    doc.set_selection(view.id, selection);
 }
 
 fn extend_line_up(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
-    let selection = doc.selection(view.id).transform(|range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         movement::move_vertically(text, range, Direction::Backward, count, Movement::Extend)
     });
-    doc.set_selection(view.id, selection);
 }
 
 fn extend_line_down(cx: &mut Context) {
     let count = cx.count();
     let (view, doc) = current!(cx.editor);
-    let text = doc.text().slice(..);
-    let selection = doc.selection(view.id).transform(|range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         movement::move_vertically(text, range, Direction::Forward, count, Movement::Extend)
     });
-    doc.set_selection(view.id, selection);
 }
 
 fn extend_line_end(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
 
-    let selection = doc.selection(view.id).transform(|range| {
-        let text = doc.text();
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         let line = text.char_to_line(range.head);
 
         let pos = line_end_char_index(&text.slice(..), line);
@@ -834,23 +787,19 @@ fn extend_line_end(cx: &mut Context) {
 
         Range::new(range.anchor, pos)
     });
-
-    doc.set_selection(view.id, selection);
 }
 
 fn extend_line_start(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
 
-    let selection = doc.selection(view.id).transform(|range| {
-        let text = doc.text();
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         let line = text.char_to_line(range.head);
 
         // adjust to start of the line
         let pos = text.line_to_char(line);
         Range::new(range.anchor, pos)
     });
-
-    doc.set_selection(view.id, selection);
 }
 
 fn select_all(cx: &mut Context) {
@@ -1043,20 +992,16 @@ fn change_selection(cx: &mut Context) {
 
 fn collapse_selection(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
-    let selection = doc
-        .selection(view.id)
-        .transform(|range| Range::new(range.head, range.head));
 
-    doc.set_selection(view.id, selection);
+    doc.selection_mut(view.id)
+        .transform(|range| Range::new(range.head, range.head));
 }
 
 fn flip_selections(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
-    let selection = doc
-        .selection(view.id)
-        .transform(|range| Range::new(range.head, range.anchor));
 
-    doc.set_selection(view.id, selection);
+    doc.selection_mut(view.id)
+        .transform(|range| Range::new(range.head, range.anchor));
 }
 
 fn enter_insert_mode(doc: &mut Document) {
@@ -1068,10 +1013,8 @@ fn insert_mode(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
     enter_insert_mode(doc);
 
-    let selection = doc
-        .selection(view.id)
+    doc.selection_mut(view.id)
         .transform(|range| Range::new(range.to(), range.from()));
-    doc.set_selection(view.id, selection);
 }
 
 // inserts at the end of each selection
@@ -1080,8 +1023,8 @@ fn append_mode(cx: &mut Context) {
     enter_insert_mode(doc);
     doc.restore_cursor = true;
 
-    let text = doc.text().slice(..);
-    let selection = doc.selection(view.id).transform(|range| {
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         Range::new(
             range.from(),
             graphemes::next_grapheme_boundary(text, range.to()), // to() + next char
@@ -1097,8 +1040,6 @@ fn append_mode(cx: &mut Context) {
         );
         doc.apply(&transaction, view.id);
     }
-
-    doc.set_selection(view.id, selection);
 }
 
 mod cmd {
@@ -1904,13 +1845,12 @@ fn append_to_line(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
     enter_insert_mode(doc);
 
-    let selection = doc.selection(view.id).transform(|range| {
-        let text = doc.text();
+    let (text, selection) = doc.text_and_mut_selection(view.id);
+    selection.transform(|range| {
         let line = text.char_to_line(range.head);
         let pos = line_end_char_index(&text.slice(..), line);
         Range::new(pos, pos)
     });
-    doc.set_selection(view.id, selection);
 }
 
 /// Sometimes when applying formatting changes we want to mark the buffer as unmodified, for
@@ -2040,14 +1980,13 @@ fn normal_mode(cx: &mut Context) {
 
     // if leaving append mode, move cursor back by 1
     if doc.restore_cursor {
-        let text = doc.text().slice(..);
-        let selection = doc.selection(view.id).transform(|range| {
+        let (text, selection) = doc.text_and_mut_selection(view.id);
+        selection.transform(|range| {
             Range::new(
                 range.from(),
                 graphemes::prev_grapheme_boundary(text, range.to()),
             )
         });
-        doc.set_selection(view.id, selection);
 
         doc.restore_cursor = false;
     }
@@ -2679,11 +2618,9 @@ pub mod insert {
     pub fn delete_word_backward(cx: &mut Context) {
         let count = cx.count();
         let (view, doc) = current!(cx.editor);
-        let text = doc.text().slice(..);
-        let selection = doc
-            .selection(view.id)
-            .transform(|range| movement::move_prev_word_start(text, range, count));
-        doc.set_selection(view.id, selection);
+
+        let (text, selection) = doc.text_and_mut_selection(view.id);
+        selection.transform(|range| movement::move_prev_word_start(text, range, count));
         delete_selection(cx)
     }
 }
diff --git a/helix-view/src/document.rs b/helix-view/src/document.rs
index 0f1f3a8f..59a1c42c 100644
--- a/helix-view/src/document.rs
+++ b/helix-view/src/document.rs
@@ -13,8 +13,8 @@ use helix_core::{
     history::History,
     line_ending::auto_detect_line_ending,
     syntax::{self, LanguageConfiguration},
-    ChangeSet, Diagnostic, LineEnding, Rope, RopeBuilder, Selection, State, Syntax, Transaction,
-    DEFAULT_LINE_ENDING,
+    ChangeSet, Diagnostic, LineEnding, Rope, RopeBuilder, RopeSlice, Selection, State, Syntax,
+    Transaction, DEFAULT_LINE_ENDING,
 };
 use helix_lsp::util::LspFormatting;
 
@@ -1000,6 +1000,22 @@ impl Document {
         &self.selections[&view_id]
     }
 
+    #[inline]
+    pub fn selection_mut(&mut self, view_id: ViewId) -> &mut Selection {
+        self.selections
+            .get_mut(&view_id)
+            .expect("No selection set with the given ViewId")
+    }
+
+    pub fn text_and_mut_selection(&mut self, view_id: ViewId) -> (RopeSlice, &mut Selection) {
+        (
+            self.text.slice(..),
+            self.selections
+                .get_mut(&view_id)
+                .expect("No selection set with the given ViewId"),
+        )
+    }
+
     pub fn relative_path(&self) -> Option<PathBuf> {
         let cwdir = std::env::current_dir().expect("couldn't determine current directory");