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 <nathanbarry474@gmail.com>
Co-authored-by: Anthony Eid <hello@anthonyeid.me>

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<Identifier, usize> = HashMap::default();
                 let mut current_rank = 0;

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::<Vec<_>>();
     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);