1 files changed, 177 insertions, 64 deletions

diff --git a/helix-lsp/src/lib.rs b/helix-lsp/src/lib.rs
index 147b381c..58e8d83d 100644
--- a/helix-lsp/src/lib.rs
+++ b/helix-lsp/src/lib.rs
@@ -59,8 +59,8 @@ pub enum OffsetEncoding {
 pub mod util {
     use super::*;
     use helix_core::line_ending::{line_end_byte_index, line_end_char_index};
+    use helix_core::{chars, RopeSlice, SmallVec};
     use helix_core::{diagnostic::NumberOrString, Range, Rope, Selection, Tendril, Transaction};
-    use helix_core::{smallvec, SmallVec};
 
     /// Converts a diagnostic in the document to [`lsp::Diagnostic`].
     ///
@@ -247,13 +247,46 @@ pub mod util {
         Some(Range::new(start, end))
     }
 
+    /// If the LS did not provide a range for the completion or the range of the
+    /// primary cursor can not be used for the secondary cursor, this function
+    /// can be used to find the completion range for a cursor
+    fn find_completion_range(text: RopeSlice, cursor: usize) -> (usize, usize) {
+        let start = cursor
+            - text
+                .chars_at(cursor)
+                .reversed()
+                .take_while(|ch| chars::char_is_word(*ch))
+                .count();
+        (start, cursor)
+    }
+    fn completion_range(
+        text: RopeSlice,
+        edit_offset: Option<(i128, i128)>,
+        cursor: usize,
+    ) -> Option<(usize, usize)> {
+        let res = match edit_offset {
+            Some((start_offset, end_offset)) => {
+                let start_offset = cursor as i128 + start_offset;
+                if start_offset < 0 {
+                    return None;
+                }
+                let end_offset = cursor as i128 + end_offset;
+                if end_offset > text.len_chars() as i128 {
+                    return None;
+                }
+                (start_offset as usize, end_offset as usize)
+            }
+            None => find_completion_range(text, cursor),
+        };
+        Some(res)
+    }
+
     /// Creates a [Transaction] from the [lsp::TextEdit] in a completion response.
     /// The transaction applies the edit to all cursors.
     pub fn generate_transaction_from_completion_edit(
         doc: &Rope,
         selection: &Selection,
-        start_offset: i128,
-        end_offset: i128,
+        edit_offset: Option<(i128, i128)>,
         new_text: String,
     ) -> Transaction {
         let replacement: Option<Tendril> = if new_text.is_empty() {
@@ -263,83 +296,163 @@ pub mod util {
         };
 
         let text = doc.slice(..);
+        let (removed_start, removed_end) =
+            completion_range(text, edit_offset, selection.primary().cursor(text))
+                .expect("transaction must be valid for primary selection");
+        let removed_text = text.slice(removed_start..removed_end);
 
-        Transaction::change_by_selection(doc, selection, |range| {
-            let cursor = range.cursor(text);
-            (
-                (cursor as i128 + start_offset) as usize,
-                (cursor as i128 + end_offset) as usize,
-                replacement.clone(),
-            )
-        })
+        let (transaction, mut selection) = Transaction::change_by_selection_ignore_overlapping(
+            doc,
+            selection,
+            |range| {
+                let cursor = range.cursor(text);
+                completion_range(text, edit_offset, cursor)
+                    .filter(|(start, end)| text.slice(start..end) == removed_text)
+                    .unwrap_or_else(|| find_completion_range(text, cursor))
+            },
+            |_, _| replacement.clone(),
+        );
+        if transaction.changes().is_empty() {
+            return transaction;
+        }
+        selection = selection.map(transaction.changes());
+        transaction.with_selection(selection)
     }
 
     /// Creates a [Transaction] from the [snippet::Snippet] in a completion response.
     /// The transaction applies the edit to all cursors.
+    #[allow(clippy::too_many_arguments)]
     pub fn generate_transaction_from_snippet(
         doc: &Rope,
         selection: &Selection,
-        start_offset: i128,
-        end_offset: i128,
+        edit_offset: Option<(i128, i128)>,
         snippet: snippet::Snippet,
         line_ending: &str,
         include_placeholder: bool,
+        tab_width: usize,
     ) -> Transaction {
         let text = doc.slice(..);
 
-        // For each cursor store offsets for the first tabstop
-        let mut cursor_tabstop_offsets = Vec::<SmallVec<[(i128, i128); 1]>>::new();
-        let transaction = Transaction::change_by_selection(doc, selection, |range| {
-            let cursor = range.cursor(text);
-            let replacement_start = (cursor as i128 + start_offset) as usize;
-            let replacement_end = (cursor as i128 + end_offset) as usize;
-            let newline_with_offset = format!(
-                "{line_ending}{blank:width$}",
-                line_ending = line_ending,
-                width = replacement_start - doc.line_to_char(doc.char_to_line(replacement_start)),
-                blank = ""
-            );
-
-            let (replacement, tabstops) =
-                snippet::render(&snippet, newline_with_offset, include_placeholder);
-
-            let replacement_len = replacement.chars().count();
-            cursor_tabstop_offsets.push(
-                tabstops
-                    .first()
-                    .unwrap_or(&smallvec![(replacement_len, replacement_len)])
-                    .iter()
-                    .map(|(from, to)| -> (i128, i128) {
-                        (
-                            *from as i128 - replacement_len as i128,
-                            *to as i128 - replacement_len as i128,
-                        )
-                    })
-                    .collect(),
-            );
-
-            (replacement_start, replacement_end, Some(replacement.into()))
-        });
+        let mut off = 0i128;
+        let mut mapped_doc = doc.clone();
+        let mut selection_tabstops: SmallVec<[_; 1]> = SmallVec::new();
+        let (removed_start, removed_end) =
+            completion_range(text, edit_offset, selection.primary().cursor(text))
+                .expect("transaction must be valid for primary selection");
+        let removed_text = text.slice(removed_start..removed_end);
 
-        // Create new selection based on the cursor tabstop from above
-        let mut cursor_tabstop_offsets_iter = cursor_tabstop_offsets.iter();
-        let selection = selection
-            .clone()
-            .map(transaction.changes())
-            .transform_iter(|range| {
-                cursor_tabstop_offsets_iter
-                    .next()
-                    .unwrap()
-                    .iter()
-                    .map(move |(from, to)| {
-                        Range::new(
-                            (range.anchor as i128 + *from) as usize,
-                            (range.anchor as i128 + *to) as usize,
-                        )
-                    })
-            });
+        let (transaction, selection) = Transaction::change_by_selection_ignore_overlapping(
+            doc,
+            selection,
+            |range| {
+                let cursor = range.cursor(text);
+                completion_range(text, edit_offset, cursor)
+                    .filter(|(start, end)| text.slice(start..end) == removed_text)
+                    .unwrap_or_else(|| find_completion_range(text, cursor))
+            },
+            |replacement_start, replacement_end| {
+                let mapped_replacement_start = (replacement_start as i128 + off) as usize;
+                let mapped_replacement_end = (replacement_end as i128 + off) as usize;
+
+                let line_idx = mapped_doc.char_to_line(mapped_replacement_start);
+                let pos_on_line = mapped_replacement_start - mapped_doc.line_to_char(line_idx);
+
+                // we only care about the actual offset here (not virtual text/softwrap)
+                // so it's ok to use the deprecated function here
+                #[allow(deprecated)]
+                let width = helix_core::visual_coords_at_pos(
+                    mapped_doc.line(line_idx),
+                    pos_on_line,
+                    tab_width,
+                )
+                .col;
+                let newline_with_offset = format!(
+                    "{line_ending}{blank:width$}",
+                    line_ending = line_ending,
+                    blank = ""
+                );
+
+                let (replacement, tabstops) =
+                    snippet::render(&snippet, &newline_with_offset, include_placeholder);
+                selection_tabstops.push((mapped_replacement_start, tabstops));
+                mapped_doc.remove(mapped_replacement_start..mapped_replacement_end);
+                mapped_doc.insert(mapped_replacement_start, &replacement);
+                off +=
+                    replacement_start as i128 - replacement_end as i128 + replacement.len() as i128;
+
+                Some(replacement)
+            },
+        );
 
-        transaction.with_selection(selection)
+        let changes = transaction.changes();
+        if changes.is_empty() {
+            return transaction;
+        }
+
+        let mut mapped_selection = SmallVec::with_capacity(selection.len());
+        let mut mapped_primary_idx = 0;
+        let primary_range = selection.primary();
+        for (range, (tabstop_anchor, tabstops)) in selection.into_iter().zip(selection_tabstops) {
+            if range == primary_range {
+                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
+                mapped_selection.push(range);
+                continue;
+            };
+
+            // expand the selection to cover the tabstop to retain the helix selection semantic
+            // the tabstop closest to the range simply replaces `head` while anchor remains in place
+            // the remaining tabstops receive their own single-width cursor
+            if range.head < range.anchor {
+                let first_tabstop = tabstop_anchor + tabstops[0].1;
+
+                // if selection is forward but was moved to the right it is
+                // contained entirely in the replacement text, just do a point
+                // selection (fallback below)
+                if range.anchor >= first_tabstop {
+                    let range = Range::new(range.anchor, first_tabstop);
+                    mapped_selection.push(range);
+                    let rem_tabstops = tabstops[1..]
+                        .iter()
+                        .map(|tabstop| Range::point(tabstop_anchor + tabstop.1));
+                    mapped_selection.extend(rem_tabstops);
+                    continue;
+                }
+            } else {
+                let last_idx = tabstops.len() - 1;
+                let last_tabstop = tabstop_anchor + tabstops[last_idx].1;
+
+                // if selection is forward but was moved to the right it is
+                // contained entirely in the replacement text, just do a point
+                // selection (fallback below)
+                if range.anchor <= last_tabstop {
+                    // we can't properly compute the the next grapheme
+                    // here because the transaction hasn't been applied yet
+                    // that is not a problem because the range gets grapheme aligned anyway
+                    // tough so just adding one will always cause head to be grapheme
+                    // aligned correctly when applied to the document
+                    let range = Range::new(range.anchor, last_tabstop + 1);
+                    mapped_selection.push(range);
+                    let rem_tabstops = tabstops[..last_idx]
+                        .iter()
+                        .map(|tabstop| Range::point(tabstop_anchor + tabstop.0));
+                    mapped_selection.extend(rem_tabstops);
+                    continue;
+                }
+            };
+
+            let tabstops = tabstops
+                .iter()
+                .map(|tabstop| Range::point(tabstop_anchor + tabstop.0));
+            mapped_selection.extend(tabstops);
+        }
+
+        transaction.with_selection(Selection::new(mapped_selection, mapped_primary_idx))
     }
 
     pub fn generate_transaction_from_edits(