From 16704212b671b5cc609d5bf45dd8e7ff8d5c1e4b Mon Sep 17 00:00:00 2001 From: Max Brunsfeld Date: Tue, 6 Jan 2026 17:52:18 -0800 Subject: [PATCH] Fix coalescing of multi-line edits --- crates/edit_prediction/src/edit_prediction.rs | 52 +++-- .../src/edit_prediction_tests.rs | 190 ++++++++++++++++++ 2 files changed, 225 insertions(+), 17 deletions(-) diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index 5c7eb9042f57b7..a8a1d803a9709e 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -425,7 +425,7 @@ struct LastEvent { new_snapshot: TextBufferSnapshot, old_file: Option>, new_file: Option>, - end_edit_anchor: Option, + edit_range: Option>, snapshot_after_last_editing_pause: Option, last_edit_time: Option, } @@ -479,7 +479,7 @@ impl LastEvent { new_snapshot: boundary_snapshot.clone(), old_file: self.old_file.clone(), new_file: self.new_file.clone(), - end_edit_anchor: self.end_edit_anchor, + edit_range: None, snapshot_after_last_editing_pause: None, last_edit_time: self.last_edit_time, }; @@ -489,7 +489,7 @@ impl LastEvent { new_snapshot: self.new_snapshot.clone(), old_file: self.old_file.clone(), new_file: self.new_file.clone(), - end_edit_anchor: self.end_edit_anchor, + edit_range: None, snapshot_after_last_editing_pause: None, last_edit_time: self.last_edit_time, }; @@ -999,10 +999,7 @@ impl EditPredictionStore { let old_file = mem::replace(&mut registered_buffer.file, new_file.clone()); let old_snapshot = mem::replace(&mut registered_buffer.snapshot, new_snapshot.clone()); - let end_edit_anchor = new_snapshot - .anchored_edits_since::(&old_snapshot.version) - .last() - .map(|(_, range)| range.end); + let edit_range = edited_range(&old_snapshot, &new_snapshot); let events = &mut project_state.events; let now = cx.background_executor().now(); @@ -1012,13 +1009,19 @@ impl EditPredictionStore { && old_snapshot.version == last_event.new_snapshot.version; let should_coalesce = is_next_snapshot_of_same_buffer - && end_edit_anchor + && edit_range .as_ref() - .zip(last_event.end_edit_anchor.as_ref()) + .zip(last_event.edit_range.as_ref()) .is_some_and(|(a, b)| { let a = a.to_point(&new_snapshot); let b = b.to_point(&new_snapshot); - a.row.abs_diff(b.row) <= CHANGE_GROUPING_LINE_SPAN + if a.start > b.end { + a.start.row.abs_diff(b.end.row) <= CHANGE_GROUPING_LINE_SPAN + } else if b.start > a.end { + b.start.row.abs_diff(a.end.row) <= CHANGE_GROUPING_LINE_SPAN + } else { + true + } }); if should_coalesce { @@ -1031,19 +1034,20 @@ impl EditPredictionStore { Some(last_event.new_snapshot.clone()); } - last_event.end_edit_anchor = end_edit_anchor; + last_event.edit_range = edit_range; last_event.new_snapshot = new_snapshot; last_event.last_edit_time = Some(now); return; } } - if events.len() + 1 >= EVENT_COUNT_MAX { - events.pop_front(); - } - if let Some(event) = project_state.last_event.take() { - events.extend(event.finalize(&project_state.license_detection_watchers, cx)); + if let Some(event) = event.finalize(&project_state.license_detection_watchers, cx) { + if events.len() + 1 >= EVENT_COUNT_MAX { + events.pop_front(); + } + events.push_back(event); + } } project_state.last_event = Some(LastEvent { @@ -1051,7 +1055,7 @@ impl EditPredictionStore { new_file, old_snapshot, new_snapshot, - end_edit_anchor, + edit_range, snapshot_after_last_editing_pause: None, last_edit_time: Some(now), }); @@ -2092,6 +2096,20 @@ impl EditPredictionStore { } } +fn edited_range( + old_snapshot: &TextBufferSnapshot, + new_snapshot: &TextBufferSnapshot, +) -> Option> { + new_snapshot + .anchored_edits_since::(&old_snapshot.version) + .fold(None, |acc, (_, range)| { + Some(match acc { + None => range, + Some(acc) => acc.start..range.end, + }) + }) +} + #[derive(Error, Debug)] #[error( "You must update to Zed version {minimum_version} or higher to continue using edit predictions." diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 40d729f1ebdced..1b98d383d306db 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -402,6 +402,196 @@ async fn test_edit_history_getter_pause_splits_last_event(cx: &mut TestAppContex ); } +#[gpui::test] +async fn test_event_grouping_line_span_coalescing(cx: &mut TestAppContext) { + let (ep_store, _requests) = init_test_with_fake_client(cx); + let fs = FakeFs::new(cx.executor()); + + // Create a file with 30 lines to test line-based coalescing + let content = (1..=30) + .map(|i| format!("Line {}\n", i)) + .collect::(); + fs.insert_tree( + "/root", + json!({ + "foo.md": content + }), + ) + .await; + let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await; + + let buffer = project + .update(cx, |project, cx| { + let path = project.find_project_path(path!("root/foo.md"), cx).unwrap(); + project.open_buffer(path, cx) + }) + .await + .unwrap(); + + ep_store.update(cx, |ep_store, cx| { + ep_store.register_buffer(&buffer, &project, cx); + }); + + // First edit: multi-line edit spanning rows 10-12 (replacing lines 11-13) + buffer.update(cx, |buffer, cx| { + let start = Point::new(10, 0).to_offset(buffer); + let end = Point::new(13, 0).to_offset(buffer); + buffer.edit(vec![(start..end, "Middle A\nMiddle B\n")], None, cx); + }); + + let events = ep_store.update(cx, |ep_store, cx| { + ep_store.edit_history_for_project(&project, cx) + }); + assert_eq!( + render_events(&events), + indoc! {" + @@ -8,9 +8,8 @@ + Line 8 + Line 9 + Line 10 + -Line 11 + -Line 12 + -Line 13 + +Middle A + +Middle B + Line 14 + Line 15 + Line 16 + "}, + "After first edit" + ); + + // Second edit: insert ABOVE the first edit's range (row 5, within 8 lines of row 10) + // This tests that coalescing considers the START of the existing range + buffer.update(cx, |buffer, cx| { + let offset = Point::new(5, 0).to_offset(buffer); + buffer.edit(vec![(offset..offset, "Above\n")], None, cx); + }); + + let events = ep_store.update(cx, |ep_store, cx| { + ep_store.edit_history_for_project(&project, cx) + }); + assert_eq!( + render_events(&events), + indoc! {" + @@ -3,14 +3,14 @@ + Line 3 + Line 4 + Line 5 + +Above + Line 6 + Line 7 + Line 8 + Line 9 + Line 10 + -Line 11 + -Line 12 + -Line 13 + +Middle A + +Middle B + Line 14 + Line 15 + Line 16 + "}, + "After inserting above (should coalesce)" + ); + + // Third edit: insert BELOW the first edit's range (row 14 in current buffer, within 8 lines of row 12) + // This tests that coalescing considers the END of the existing range + buffer.update(cx, |buffer, cx| { + let offset = Point::new(14, 0).to_offset(buffer); + buffer.edit(vec![(offset..offset, "Below\n")], None, cx); + }); + + let events = ep_store.update(cx, |ep_store, cx| { + ep_store.edit_history_for_project(&project, cx) + }); + assert_eq!( + render_events(&events), + indoc! {" + @@ -3,15 +3,16 @@ + Line 3 + Line 4 + Line 5 + +Above + Line 6 + Line 7 + Line 8 + Line 9 + Line 10 + -Line 11 + -Line 12 + -Line 13 + +Middle A + +Middle B + Line 14 + +Below + Line 15 + Line 16 + Line 17 + "}, + "After inserting below (should coalesce)" + ); + + // Fourth edit: insert FAR BELOW (row 25, beyond 8 lines from the current range end ~row 15) + // This should NOT coalesce - creates a new event + buffer.update(cx, |buffer, cx| { + let offset = Point::new(25, 0).to_offset(buffer); + buffer.edit(vec![(offset..offset, "Far below\n")], None, cx); + }); + + let events = ep_store.update(cx, |ep_store, cx| { + ep_store.edit_history_for_project(&project, cx) + }); + assert_eq!( + render_events(&events), + indoc! {" + @@ -3,15 +3,16 @@ + Line 3 + Line 4 + Line 5 + +Above + Line 6 + Line 7 + Line 8 + Line 9 + Line 10 + -Line 11 + -Line 12 + -Line 13 + +Middle A + +Middle B + Line 14 + +Below + Line 15 + Line 16 + Line 17 + + --- + @@ -23,6 +23,7 @@ + Line 22 + Line 23 + Line 24 + +Far below + Line 25 + Line 26 + Line 27 + "}, + "After inserting far below (should NOT coalesce)" + ); +} + +fn render_events(events: &[StoredEvent]) -> String { + events + .iter() + .map(|e| { + let zeta_prompt::Event::BufferChange { diff, .. } = e.event.as_ref(); + diff.as_str() + }) + .collect::>() + .join("\n---\n") +} + #[gpui::test] async fn test_empty_prediction(cx: &mut TestAppContext) { let (ep_store, mut requests) = init_test_with_fake_client(cx);