From bf19f03f51e08e668c0e4eb82df4833b10b19bcd Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Sat, 24 Jan 2026 18:07:59 -0800
Subject: [PATCH] Apply common prefix/suffix stripping to zeta2 and mercury
 (#47530)

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 ...
---
 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(-)

diff --git a/crates/edit_prediction/src/mercury.rs b/crates/edit_prediction/src/mercury.rs
index 8a6c4f92d49c663b8baeb17c5f2ebd99230f2c51..09d301c5fa2e7c0edf964c626d3d40d9764c33cc 100644
--- a/crates/edit_prediction/src/mercury.rs
+++ b/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,
                 );
             }
 
diff --git a/crates/edit_prediction/src/zeta1.rs b/crates/edit_prediction/src/zeta1.rs
index cb3897a90f2bd510f24d97d3465cb4ba8ee358b3..216074aae15d75aea9f43592d64097d2153e9e6a 100644
--- a/crates/edit_prediction/src/zeta1.rs
+++ b/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() {
diff --git a/crates/edit_prediction/src/zeta2.rs b/crates/edit_prediction/src/zeta2.rs
index 9eefe1e38e1fcf17db575d92a62dc37bbcf5eb1e..3cdd0ff8a5751bd268e5f01053efc86c6a8fab9a 100644
--- a/crates/edit_prediction/src/zeta2.rs
+++ b/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((