Apply common prefix/suffix stripping to zeta2 and mercury (#47530)

Max Brunsfeld created 2 weeks ago

Fixes an issue where edits would be shown like this:

<img width="1262" height="1072" alt="image"
src="https://github.com/user-attachments/assets/f3ec865a-4bf2-423c-8465-62a54a21a882"
/>

When they should be like this:

<img width="1260" height="388" alt="image"
src="https://github.com/user-attachments/assets/ac555a25-f222-48aa-b922-0ae5b9f4e260"
/>


Release Notes:

- N/A *or* Added/Fixed/Improved ...

Change summary

crates/edit_prediction/src/mercury.rs | 18 ++++++------------
crates/edit_prediction/src/zeta1.rs   | 18 ++++++++----------
crates/edit_prediction/src/zeta2.rs   | 17 +++++++----------
3 files changed, 21 insertions(+), 32 deletions(-)

Detailed changes

crates/edit_prediction/src/mercury.rs

@@ -1,7 +1,7 @@
 use crate::{
     DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionModelInput,
     EditPredictionStartedDebugEvent, open_ai_response::text_from_response,
-    prediction::EditPredictionResult,
+    prediction::EditPredictionResult, zeta1::compute_edits,
 };
 use anyhow::{Context as _, Result};
 use futures::AsyncReadExt as _;
@@ -184,17 +184,11 @@ impl Mercury {
                 let old_text = snapshot
                     .text_for_range(editable_offset_range.clone())
                     .collect::<String>();
-                edits.extend(
-                    language::text_diff(&old_text, &response_str)
-                        .into_iter()
-                        .map(|(range, text)| {
-                            (
-                                snapshot.anchor_after(editable_offset_range.start + range.start)
-                                    ..snapshot
-                                        .anchor_before(editable_offset_range.start + range.end),
-                                text,
-                            )
-                        }),
+                edits = compute_edits(
+                    old_text,
+                    &response_str,
+                    editable_offset_range.start,
+                    &snapshot,
                 );
             }
 

crates/edit_prediction/src/zeta1.rs

@@ -347,20 +347,18 @@ pub fn compute_edits(
     text_diff(&old_text, new_text)
         .into_iter()
         .map(|(mut old_range, new_text)| {
-            old_range.start += offset;
-            old_range.end += offset;
-
-            let prefix_len = common_prefix(
-                snapshot.chars_for_range(old_range.clone()),
-                new_text.chars(),
-            );
-            old_range.start += prefix_len;
+            let old_slice = &old_text[old_range.clone()];
 
+            let prefix_len = common_prefix(old_slice.chars(), new_text.chars());
             let suffix_len = common_prefix(
-                snapshot.reversed_chars_for_range(old_range.clone()),
+                old_slice[prefix_len..].chars().rev(),
                 new_text[prefix_len..].chars().rev(),
             );
-            old_range.end = old_range.end.saturating_sub(suffix_len);
+
+            old_range.start += offset;
+            old_range.end += offset;
+            old_range.start += prefix_len;
+            old_range.end -= suffix_len;
 
             let new_text = new_text[prefix_len..new_text.len() - suffix_len].into();
             let range = if old_range.is_empty() {

crates/edit_prediction/src/zeta2.rs

@@ -1,4 +1,5 @@
 use crate::prediction::EditPredictionResult;
+use crate::zeta1::compute_edits;
 use crate::{
     CurrentEditPrediction, DebugEvent, EDIT_PREDICTIONS_MODEL_ID, EditPredictionFinishedDebugEvent,
     EditPredictionId, EditPredictionModelInput, EditPredictionStartedDebugEvent,
@@ -168,16 +169,12 @@ pub fn request_prediction_with_zeta2(
                 old_text.push('\n');
             }
 
-            let edits: Vec<_> = language::text_diff(&old_text, &output_text)
-                .into_iter()
-                .map(|(range, text)| {
-                    (
-                        snapshot.anchor_after(editable_offset_range.start + range.start)
-                            ..snapshot.anchor_before(editable_offset_range.start + range.end),
-                        text,
-                    )
-                })
-                .collect();
+            let edits = compute_edits(
+                old_text,
+                &output_text,
+                editable_offset_range.start,
+                &snapshot,
+            );
 
             anyhow::Ok((
                 Some((