From 02275294af3cdb7d84ef9ff5561a21ade36879bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Miguel=20Raz=20Guzm=C3=A1n=20Macedo?= Date: Mon, 8 Jun 2026 16:43:48 -0600 Subject: [PATCH] Limit top context identifiers and diagnostics (#58870) - Replace full sort and truncate with a budgeted top-N approach for context identifiers using truncate_to_bottom_n_sorted_by. - Introduce MAX_DIAGNOSTICS = 10 and prune diagnostics to this limit with select_nth_unstable_by_key, then sort the remainder. Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content adheres to Zed's UI standards ([UX/UI](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) and [icon](https://github.com/zed-industries/zed/blob/main/crates/icons/README.md) guidelines) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Release Notes: - Removed small allocations on the edit_predictions critical path by using select_nth and sort_unstable Co-authored-by: Nathan Barry Co-authored-by: Anthony Eid --- .../src/edit_prediction_context.rs | 9 +++++++-- crates/zeta_prompt/src/zeta_prompt.rs | 15 ++++++++++++--- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/crates/edit_prediction_context/src/edit_prediction_context.rs b/crates/edit_prediction_context/src/edit_prediction_context.rs index 2b9bb38015..0bcd63cdcf 100644 --- a/crates/edit_prediction_context/src/edit_prediction_context.rs +++ b/crates/edit_prediction_context/src/edit_prediction_context.rs @@ -260,8 +260,13 @@ impl RelatedExcerptStore { (id, distance) }) .collect(); - identifiers_with_distance.sort_by_key(|(_, distance)| *distance); - identifiers_with_distance.truncate(MAX_CONTEXT_IDENTIFIER_COUNT); + // Only the closest `MAX_CONTEXT_IDENTIFIER_COUNT` identifiers are + // used below, so select that prefix instead of fully sorting. + util::truncate_to_bottom_n_sorted_by( + &mut identifiers_with_distance, + MAX_CONTEXT_IDENTIFIER_COUNT, + &|(_, a), (_, b)| a.cmp(b), + ); let mut cursor_distances: HashMap = HashMap::default(); let mut current_rank = 0; diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 335eb85f29..f66ee67f46 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -920,13 +920,22 @@ pub fn format_active_buffer_diagnostics_with_budget( return String::new(); } + const MAX_DIAGNOSTICS: usize = 10; + let mut diagnostic_indices = (0..diagnostics.len()).collect::>(); if let Some(cursor_buffer_row) = cursor_buffer_row { - diagnostic_indices.sort_by_key(|index| { + let distance = |index: &usize| { let range = &diagnostics[*index].snippet_buffer_row_range; u32::abs_diff(cursor_buffer_row, range.start) + u32::abs_diff(cursor_buffer_row, range.end) - }); + }; + // Only the closest `MAX_DIAGNOSTICS` are rendered below, so select that + // prefix instead of fully sorting every diagnostic. + if diagnostic_indices.len() > MAX_DIAGNOSTICS { + diagnostic_indices.select_nth_unstable_by_key(MAX_DIAGNOSTICS, &distance); + diagnostic_indices.truncate(MAX_DIAGNOSTICS); + } + diagnostic_indices.sort_unstable_by_key(&distance); } let mut output = format!("{}diagnostics\n", seed_coder::FILE_MARKER); @@ -937,7 +946,7 @@ pub fn format_active_buffer_diagnostics_with_budget( let mut used_tokens = header_tokens; let mut included_diagnostics = 0; - for diagnostic_index in diagnostic_indices.into_iter().take(10) { + for diagnostic_index in diagnostic_indices.into_iter().take(MAX_DIAGNOSTICS) { let diagnostic = &diagnostics[diagnostic_index]; let snippet = clamp_text_to_token_count(&diagnostic.snippet, 256);