edit prediction: Fix crash in `highlight_text` (#23766)

Bennet Bo Fenner and Antonio created 11 months ago

This fixes the panics we we're seeing in `EditPreview::highlight_edits`.
The reason for this was that we were interpolating edits incorrectly.

Here's an example:

```rust
let a = 0; // existing code
let c = 2; // suggested by edit prediction
```
The edits would look like this: `[(Point(1, 0)..Point(1, 0), "let c =
2;"]`

Now i type:
```rust
let a = 0; // existing code
let b = 1; // added this line
let c = 2; // suggested by edit prediction
```

Before this change, the `interpolate` function would allow insertions
before the edit prediction edits, the anchors will move to the next
line.
The edits would look now like this: `[(Point(2, 0)..Point(2, 0), "let c
= 2;"]`

However, now we end up with a call to `EditPreview::highlight_edits`,
with the following parameters:
- current_snapshot: 
  ```rust
  let a = 0;
  let b = 1;
  ```
- edits: `[(Point(2, 0)..Point(2, 0), "let c = 2;"]`
- applied_edits_snapshot:
  ```rust
  let a = 0;
  let c = 2;
  ```

And here you can see the issue, applying the `edits` to the
`current_snapshot` should always end up re-creating the text that is
present in the `applied_edits_snapshot`. That is not the case here
though, meaning that the offsets in the new buffer are not correct,
which can either lead to a confusing popup or a crash if the suggestion
is at the end of the file.

Here's a real world example (edit prediction is ONLY suggesting to
delete a new line):

<img width="487" alt="Screenshot 2025-01-27 at 13 05 26"
src="https://github.com/user-attachments/assets/a0a8064e-8cfa-48b2-9f1c-efc2d0d9d7d4"
/>

We fixed this by only allowing interpolation if the user is editing
after all the edit predictions OR if the user edit is a subset of the
model suggestion.



Co-Authored-by: Antonio <[e-mail]>

Release Notes:

- N/A

Co-authored-by: Antonio <[e-mail]>

Change summary

crates/zeta/src/zeta.rs | 103 +++++++++++++++++-------------------------
1 file changed, 42 insertions(+), 61 deletions(-)

Detailed changes

crates/zeta/src/zeta.rs

@@ -103,23 +103,23 @@ fn interpolate(
 ) -> Option<Vec<(Range<Anchor>, String)>> {
     let mut edits = Vec::new();
 
-    let mut user_edits = new_snapshot
-        .edits_since::<usize>(&old_snapshot.version)
-        .peekable();
-    for (model_old_range, model_new_text) in current_edits.iter() {
-        let model_offset_range = model_old_range.to_offset(old_snapshot);
-        while let Some(next_user_edit) = user_edits.peek() {
-            if next_user_edit.old.end < model_offset_range.start {
-                user_edits.next();
+    let mut model_edits = current_edits.into_iter().peekable();
+    for user_edit in new_snapshot.edits_since::<usize>(&old_snapshot.version) {
+        while let Some((model_old_range, _)) = model_edits.peek() {
+            let model_old_range = model_old_range.to_offset(old_snapshot);
+            if !model_old_range.is_empty() {
+                return None;
+            } else if model_old_range.end < user_edit.old.start {
+                let (model_old_range, model_new_text) = model_edits.next().unwrap();
+                edits.push((model_old_range.clone(), model_new_text.clone()));
             } else {
                 break;
             }
         }
 
-        if let Some(user_edit) = user_edits.peek() {
-            if user_edit.old.start > model_offset_range.end {
-                edits.push((model_old_range.clone(), model_new_text.clone()));
-            } else if user_edit.old == model_offset_range {
+        if let Some((model_old_range, model_new_text)) = model_edits.peek() {
+            let model_old_offset_range = model_old_range.to_offset(old_snapshot);
+            if user_edit.old == model_old_offset_range {
                 let user_new_text = new_snapshot
                     .text_for_range(user_edit.new.clone())
                     .collect::<String>();
@@ -128,20 +128,26 @@ fn interpolate(
                     if !model_suffix.is_empty() {
                         edits.push((
                             new_snapshot.anchor_after(user_edit.new.end)
-                                ..new_snapshot.anchor_before(user_edit.new.end),
-                            model_suffix.into(),
+                                ..new_snapshot.anchor_after(user_edit.new.end),
+                            model_suffix.to_string(),
                         ));
                     }
 
-                    user_edits.next();
-                } else {
-                    return None;
+                    model_edits.next();
+                    continue;
                 }
-            } else {
-                return None;
             }
-        } else {
+        }
+
+        return None;
+    }
+
+    for (model_old_range, model_new_text) in model_edits {
+        let model_old_offset_range = model_old_range.to_offset(old_snapshot);
+        if model_old_offset_range.is_empty() {
             edits.push((model_old_range.clone(), model_new_text.clone()));
+        } else {
+            return None;
         }
     }
 
@@ -1277,10 +1283,10 @@ mod tests {
 
     #[gpui::test]
     async fn test_inline_completion_basic_interpolation(cx: &mut TestAppContext) {
-        let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
+        let buffer = cx.new(|cx| Buffer::local("Lo ips dolor", cx));
         let edits: Arc<[(Range<Anchor>, String)]> = cx.update(|cx| {
             to_completion_edits(
-                [(2..5, "REM".to_string()), (9..11, "".to_string())],
+                [(2..2, "REM".to_string()), (6..6, "UM".to_string())],
                 &buffer,
                 cx,
             )
@@ -1314,80 +1320,55 @@ mod tests {
                     &buffer,
                     cx
                 ),
-                vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
+                vec![(2..2, "REM".to_string()), (6..6, "UM".to_string())]
             );
 
-            buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
+            buffer.update(cx, |buffer, cx| buffer.edit([(2..2, "R")], None, cx));
             assert_eq!(
                 from_completion_edits(
                     &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                     &buffer,
                     cx
                 ),
-                vec![(2..2, "REM".to_string()), (6..8, "".to_string())]
+                vec![(3..3, "EM".to_string()), (7..7, "UM".to_string())]
             );
 
-            buffer.update(cx, |buffer, cx| buffer.undo(cx));
+            buffer.update(cx, |buffer, cx| buffer.edit([(7..7, "U")], None, cx));
             assert_eq!(
                 from_completion_edits(
                     &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                     &buffer,
                     cx
                 ),
-                vec![(2..5, "REM".to_string()), (9..11, "".to_string())]
+                vec![(3..3, "EM".to_string()), (8..8, "M".to_string())]
             );
 
-            buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
-            assert_eq!(
-                from_completion_edits(
-                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
-                    &buffer,
-                    cx
-                ),
-                vec![(3..3, "EM".to_string()), (7..9, "".to_string())]
-            );
-
-            buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
-            assert_eq!(
-                from_completion_edits(
-                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
-                    &buffer,
-                    cx
-                ),
-                vec![(4..4, "M".to_string()), (8..10, "".to_string())]
-            );
-
-            buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
-            assert_eq!(
-                from_completion_edits(
-                    &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
-                    &buffer,
-                    cx
-                ),
-                vec![(9..11, "".to_string())]
-            );
+            // Editing after all the model edits should always keep the inline completion
+            buffer.update(cx, |buffer, cx| buffer.edit([(10..10, "X")], None, cx));
+            assert_eq!(completion.interpolate(&buffer.read(cx).snapshot()), None);
 
-            buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
+            buffer.update(cx, |buffer, cx| buffer.undo(cx));
             assert_eq!(
                 from_completion_edits(
                     &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                     &buffer,
                     cx
                 ),
-                vec![(4..4, "M".to_string()), (8..10, "".to_string())]
+                vec![(3..3, "EM".to_string()), (8..8, "M".to_string())]
             );
 
-            buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
+            buffer.update(cx, |buffer, cx| buffer.edit([(8..8, "M")], None, cx));
             assert_eq!(
                 from_completion_edits(
                     &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                     &buffer,
                     cx
                 ),
-                vec![(4..4, "M".to_string())]
+                vec![(3..3, "EM".to_string())]
             );
 
-            buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
+            // Editing before the model edits should always remove the inline completion
+            buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "I")], None, cx));
             assert_eq!(completion.interpolate(&buffer.read(cx).snapshot()), None);
         })
     }