1//! Zeta2 prompt planning and generation code shared with cloud.
2
3use anyhow::{Context as _, Result, anyhow};
4use cloud_llm_client::predict_edits_v3::{
5 self, Excerpt, Line, Point, PromptFormat, ReferencedDeclaration,
6};
7use indoc::indoc;
8use ordered_float::OrderedFloat;
9use rustc_hash::{FxHashMap, FxHashSet};
10use serde::Serialize;
11use std::cmp;
12use std::fmt::Write;
13use std::sync::Arc;
14use std::{cmp::Reverse, collections::BinaryHeap, ops::Range, path::Path};
15use strum::{EnumIter, IntoEnumIterator};
16
17pub const DEFAULT_MAX_PROMPT_BYTES: usize = 10 * 1024;
18
19pub const CURSOR_MARKER: &str = "<|user_cursor|>";
20/// NOTE: Differs from zed version of constant - includes a newline
21pub const EDITABLE_REGION_START_MARKER_WITH_NEWLINE: &str = "<|editable_region_start|>\n";
22/// NOTE: Differs from zed version of constant - includes a newline
23pub const EDITABLE_REGION_END_MARKER_WITH_NEWLINE: &str = "<|editable_region_end|>\n";
24
25// TODO: use constants for markers?
26const MARKED_EXCERPT_INSTRUCTIONS: &str = indoc! {"
27 You are a code completion assistant and your task is to analyze user edits and then rewrite an excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking into account the cursor location.
28
29 The excerpt to edit will be wrapped in markers <|editable_region_start|> and <|editable_region_end|>. The cursor position is marked with <|user_cursor|>. Please respond with edited code for that region.
30
31 Other code is provided for context, and `…` indicates when code has been skipped.
32
33 # Edit History:
34
35"};
36
37const LABELED_SECTIONS_INSTRUCTIONS: &str = indoc! {r#"
38 You are a code completion assistant and your task is to analyze user edits, and suggest an edit to one of the provided sections of code.
39
40 Sections of code are grouped by file and then labeled by `<|section_N|>` (e.g `<|section_8|>`).
41
42 The cursor position is marked with `<|user_cursor|>` and it will appear within a special section labeled `<|current_section|>`. Prefer editing the current section until no more changes are needed within it.
43
44 Respond ONLY with the name of the section to edit on a single line, followed by all of the code that should replace that section. For example:
45
46 <|current_section|>
47 for i in 0..16 {
48 println!("{i}");
49 }
50
51 # Edit History:
52
53"#};
54
55const NUMBERED_LINES_INSTRUCTIONS: &str = indoc! {r#"
56 # Instructions
57
58 You are a code completion assistant helping a programmer finish their work. Your task is to:
59
60 1. Analyze the edit history to understand what the programmer is trying to achieve
61 2. Identify any incomplete refactoring or changes that need to be finished
62 3. Make the remaining edits that a human programmer would logically make next
63 4. Apply systematic changes consistently across the entire codebase - if you see a pattern starting, complete it everywhere.
64
65 Focus on:
66 - Understanding the intent behind the changes (e.g., improving error handling, refactoring APIs, fixing bugs)
67 - Completing any partially-applied changes across the codebase
68 - Ensuring consistency with the programming style and patterns already established
69 - Making edits that maintain or improve code quality
70 - If the programmer started refactoring one instance of a pattern, find and update ALL similar instances
71 - Don't write a lot of code if you're not sure what to do
72
73 Rules:
74 - Do not just mechanically apply patterns - reason about what changes make sense given the context and the programmer's apparent goals.
75 - Do not just fix syntax errors - look for the broader refactoring pattern and apply it systematically throughout the code.
76 - Write the edits in the unified diff format as shown in the example.
77
78 # Example output:
79
80 ```
81 --- a/src/myapp/cli.py
82 +++ b/src/myapp/cli.py
83 @@ -1,3 +1,3 @@
84 -
85 -
86 -import sys
87 +import json
88 ```
89
90 # Edit History:
91
92"#};
93
94const UNIFIED_DIFF_REMINDER: &str = indoc! {"
95 ---
96
97 Please analyze the edit history and the files, then provide the unified diff for your predicted edits.
98 Do not include the cursor marker in your output.
99 If you're editing multiple files, be sure to reflect filename in the hunk's header.
100"};
101
102pub fn build_prompt(
103 request: &predict_edits_v3::PredictEditsRequest,
104) -> Result<(String, SectionLabels)> {
105 let mut insertions = match request.prompt_format {
106 PromptFormat::MarkedExcerpt => vec![
107 (
108 Point {
109 line: request.excerpt_line_range.start,
110 column: 0,
111 },
112 EDITABLE_REGION_START_MARKER_WITH_NEWLINE,
113 ),
114 (request.cursor_point, CURSOR_MARKER),
115 (
116 Point {
117 line: request.excerpt_line_range.end,
118 column: 0,
119 },
120 EDITABLE_REGION_END_MARKER_WITH_NEWLINE,
121 ),
122 ],
123 PromptFormat::LabeledSections => vec![(request.cursor_point, CURSOR_MARKER)],
124 PromptFormat::NumLinesUniDiff => {
125 vec![(request.cursor_point, CURSOR_MARKER)]
126 }
127 PromptFormat::OnlySnippets => vec![],
128 };
129
130 let mut prompt = match request.prompt_format {
131 PromptFormat::MarkedExcerpt => MARKED_EXCERPT_INSTRUCTIONS.to_string(),
132 PromptFormat::LabeledSections => LABELED_SECTIONS_INSTRUCTIONS.to_string(),
133 PromptFormat::NumLinesUniDiff => NUMBERED_LINES_INSTRUCTIONS.to_string(),
134 // only intended for use via zeta_cli
135 PromptFormat::OnlySnippets => String::new(),
136 };
137
138 if request.events.is_empty() {
139 prompt.push_str("(No edit history)\n\n");
140 } else {
141 prompt.push_str(
142 "The following are the latest edits made by the user, from earlier to later.\n\n",
143 );
144 push_events(&mut prompt, &request.events);
145 }
146
147 if request.prompt_format == PromptFormat::NumLinesUniDiff {
148 if request.referenced_declarations.is_empty() {
149 prompt.push_str(indoc! {"
150 # File under the cursor:
151
152 The cursor marker <|user_cursor|> indicates the current user cursor position.
153 The file is in current state, edits from edit history have been applied.
154 We prepend line numbers (e.g., `123|<actual line>`); they are not part of the file.
155
156 "});
157 } else {
158 // Note: This hasn't been trained on yet
159 prompt.push_str(indoc! {"
160 # Code Excerpts:
161
162 The cursor marker <|user_cursor|> indicates the current user cursor position.
163 Other excerpts of code from the project have been included as context based on their similarity to the code under the cursor.
164 Context excerpts are not guaranteed to be relevant, so use your own judgement.
165 Files are in their current state, edits from edit history have been applied.
166 We prepend line numbers (e.g., `123|<actual line>`); they are not part of the file.
167
168 "});
169 }
170 } else {
171 prompt.push_str("\n## Code\n\n");
172 }
173
174 let mut section_labels = Default::default();
175
176 if !request.referenced_declarations.is_empty() || !request.signatures.is_empty() {
177 let syntax_based_prompt = SyntaxBasedPrompt::populate(request)?;
178 section_labels = syntax_based_prompt.write(&mut insertions, &mut prompt)?;
179 } else {
180 if request.prompt_format == PromptFormat::LabeledSections {
181 anyhow::bail!("PromptFormat::LabeledSections cannot be used with ContextMode::Llm");
182 }
183
184 for related_file in &request.included_files {
185 write_codeblock(
186 &related_file.path,
187 &related_file.excerpts,
188 if related_file.path == request.excerpt_path {
189 &insertions
190 } else {
191 &[]
192 },
193 related_file.max_row,
194 request.prompt_format == PromptFormat::NumLinesUniDiff,
195 &mut prompt,
196 );
197 }
198 }
199
200 if request.prompt_format == PromptFormat::NumLinesUniDiff {
201 prompt.push_str(UNIFIED_DIFF_REMINDER);
202 }
203
204 Ok((prompt, section_labels))
205}
206
207pub fn write_codeblock<'a>(
208 path: &Path,
209 excerpts: impl IntoIterator<Item = &'a Excerpt>,
210 sorted_insertions: &[(Point, &str)],
211 file_line_count: Line,
212 include_line_numbers: bool,
213 output: &'a mut String,
214) {
215 writeln!(output, "`````{}", path.display()).unwrap();
216 write_excerpts(
217 excerpts,
218 sorted_insertions,
219 file_line_count,
220 include_line_numbers,
221 output,
222 );
223 write!(output, "`````\n\n").unwrap();
224}
225
226pub fn write_excerpts<'a>(
227 excerpts: impl IntoIterator<Item = &'a Excerpt>,
228 sorted_insertions: &[(Point, &str)],
229 file_line_count: Line,
230 include_line_numbers: bool,
231 output: &mut String,
232) {
233 let mut current_row = Line(0);
234 let mut sorted_insertions = sorted_insertions.iter().peekable();
235
236 for excerpt in excerpts {
237 if excerpt.start_line > current_row {
238 writeln!(output, "…").unwrap();
239 }
240 if excerpt.text.is_empty() {
241 return;
242 }
243
244 current_row = excerpt.start_line;
245
246 for mut line in excerpt.text.lines() {
247 if include_line_numbers {
248 write!(output, "{}|", current_row.0 + 1).unwrap();
249 }
250
251 while let Some((insertion_location, insertion_marker)) = sorted_insertions.peek() {
252 match current_row.cmp(&insertion_location.line) {
253 cmp::Ordering::Equal => {
254 let (prefix, suffix) = line.split_at(insertion_location.column as usize);
255 output.push_str(prefix);
256 output.push_str(insertion_marker);
257 line = suffix;
258 sorted_insertions.next();
259 }
260 cmp::Ordering::Less => break,
261 cmp::Ordering::Greater => {
262 sorted_insertions.next();
263 break;
264 }
265 }
266 }
267 output.push_str(line);
268 output.push('\n');
269 current_row.0 += 1;
270 }
271 }
272
273 if current_row < file_line_count {
274 writeln!(output, "…").unwrap();
275 }
276}
277
278fn push_events(output: &mut String, events: &[predict_edits_v3::Event]) {
279 if events.is_empty() {
280 return;
281 };
282
283 writeln!(output, "`````diff").unwrap();
284 for event in events {
285 writeln!(output, "{}", event).unwrap();
286 }
287 writeln!(output, "`````\n").unwrap();
288}
289
290pub struct SyntaxBasedPrompt<'a> {
291 request: &'a predict_edits_v3::PredictEditsRequest,
292 /// Snippets to include in the prompt. These may overlap - they are merged / deduplicated in
293 /// `to_prompt_string`.
294 snippets: Vec<PlannedSnippet<'a>>,
295 budget_used: usize,
296}
297
298#[derive(Clone, Debug)]
299pub struct PlannedSnippet<'a> {
300 path: Arc<Path>,
301 range: Range<Line>,
302 text: &'a str,
303 // TODO: Indicate this in the output
304 #[allow(dead_code)]
305 text_is_truncated: bool,
306}
307
308#[derive(EnumIter, Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
309pub enum DeclarationStyle {
310 Signature,
311 Declaration,
312}
313
314#[derive(Default, Clone, Debug, Serialize)]
315pub struct SectionLabels {
316 pub excerpt_index: usize,
317 pub section_ranges: Vec<(Arc<Path>, Range<Line>)>,
318}
319
320impl<'a> SyntaxBasedPrompt<'a> {
321 /// Greedy one-pass knapsack algorithm to populate the prompt plan. Does the following:
322 ///
323 /// Initializes a priority queue by populating it with each snippet, finding the
324 /// DeclarationStyle that minimizes `score_density = score / snippet.range(style).len()`. When a
325 /// "signature" snippet is popped, insert an entry for the "declaration" variant that reflects
326 /// the cost of upgrade.
327 ///
328 /// TODO: Implement an early halting condition. One option might be to have another priority
329 /// queue where the score is the size, and update it accordingly. Another option might be to
330 /// have some simpler heuristic like bailing after N failed insertions, or based on how much
331 /// budget is left.
332 ///
333 /// TODO: Has the current known sources of imprecision:
334 ///
335 /// * Does not consider snippet overlap when ranking. For example, it might add a field to the
336 /// plan even though the containing struct is already included.
337 ///
338 /// * Does not consider cost of signatures when ranking snippets - this is tricky since
339 /// signatures may be shared by multiple snippets.
340 ///
341 /// * Does not include file paths / other text when considering max_bytes.
342 pub fn populate(request: &'a predict_edits_v3::PredictEditsRequest) -> Result<Self> {
343 let mut this = Self {
344 request,
345 snippets: Vec::new(),
346 budget_used: request.excerpt.len(),
347 };
348 let mut included_parents = FxHashSet::default();
349 let additional_parents = this.additional_parent_signatures(
350 &request.excerpt_path,
351 request.excerpt_parent,
352 &included_parents,
353 )?;
354 this.add_parents(&mut included_parents, additional_parents);
355
356 let max_bytes = request.prompt_max_bytes.unwrap_or(DEFAULT_MAX_PROMPT_BYTES);
357
358 if this.budget_used > max_bytes {
359 return Err(anyhow!(
360 "Excerpt + signatures size of {} already exceeds budget of {}",
361 this.budget_used,
362 max_bytes
363 ));
364 }
365
366 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
367 struct QueueEntry {
368 score_density: OrderedFloat<f32>,
369 declaration_index: usize,
370 style: DeclarationStyle,
371 }
372
373 // Initialize priority queue with the best score for each snippet.
374 let mut queue: BinaryHeap<QueueEntry> = BinaryHeap::new();
375 for (declaration_index, declaration) in request.referenced_declarations.iter().enumerate() {
376 let (style, score_density) = DeclarationStyle::iter()
377 .map(|style| {
378 (
379 style,
380 OrderedFloat(declaration_score_density(&declaration, style)),
381 )
382 })
383 .max_by_key(|(_, score_density)| *score_density)
384 .unwrap();
385 queue.push(QueueEntry {
386 score_density,
387 declaration_index,
388 style,
389 });
390 }
391
392 // Knapsack selection loop
393 while let Some(queue_entry) = queue.pop() {
394 let Some(declaration) = request
395 .referenced_declarations
396 .get(queue_entry.declaration_index)
397 else {
398 return Err(anyhow!(
399 "Invalid declaration index {}",
400 queue_entry.declaration_index
401 ));
402 };
403
404 let mut additional_bytes = declaration_size(declaration, queue_entry.style);
405 if this.budget_used + additional_bytes > max_bytes {
406 continue;
407 }
408
409 let additional_parents = this.additional_parent_signatures(
410 &declaration.path,
411 declaration.parent_index,
412 &mut included_parents,
413 )?;
414 additional_bytes += additional_parents
415 .iter()
416 .map(|(_, snippet)| snippet.text.len())
417 .sum::<usize>();
418 if this.budget_used + additional_bytes > max_bytes {
419 continue;
420 }
421
422 this.budget_used += additional_bytes;
423 this.add_parents(&mut included_parents, additional_parents);
424 let planned_snippet = match queue_entry.style {
425 DeclarationStyle::Signature => {
426 let Some(text) = declaration.text.get(declaration.signature_range.clone())
427 else {
428 return Err(anyhow!(
429 "Invalid declaration signature_range {:?} with text.len() = {}",
430 declaration.signature_range,
431 declaration.text.len()
432 ));
433 };
434 let signature_start_line = declaration.range.start
435 + Line(
436 declaration.text[..declaration.signature_range.start]
437 .lines()
438 .count() as u32,
439 );
440 let signature_end_line = signature_start_line
441 + Line(
442 declaration.text
443 [declaration.signature_range.start..declaration.signature_range.end]
444 .lines()
445 .count() as u32,
446 );
447 let range = signature_start_line..signature_end_line;
448
449 PlannedSnippet {
450 path: declaration.path.clone(),
451 range,
452 text,
453 text_is_truncated: declaration.text_is_truncated,
454 }
455 }
456 DeclarationStyle::Declaration => PlannedSnippet {
457 path: declaration.path.clone(),
458 range: declaration.range.clone(),
459 text: &declaration.text,
460 text_is_truncated: declaration.text_is_truncated,
461 },
462 };
463 this.snippets.push(planned_snippet);
464
465 // When a Signature is consumed, insert an entry for Definition style.
466 if queue_entry.style == DeclarationStyle::Signature {
467 let signature_size = declaration_size(&declaration, DeclarationStyle::Signature);
468 let declaration_size =
469 declaration_size(&declaration, DeclarationStyle::Declaration);
470 let signature_score = declaration_score(&declaration, DeclarationStyle::Signature);
471 let declaration_score =
472 declaration_score(&declaration, DeclarationStyle::Declaration);
473
474 let score_diff = declaration_score - signature_score;
475 let size_diff = declaration_size.saturating_sub(signature_size);
476 if score_diff > 0.0001 && size_diff > 0 {
477 queue.push(QueueEntry {
478 declaration_index: queue_entry.declaration_index,
479 score_density: OrderedFloat(score_diff / (size_diff as f32)),
480 style: DeclarationStyle::Declaration,
481 });
482 }
483 }
484 }
485
486 anyhow::Ok(this)
487 }
488
489 fn add_parents(
490 &mut self,
491 included_parents: &mut FxHashSet<usize>,
492 snippets: Vec<(usize, PlannedSnippet<'a>)>,
493 ) {
494 for (parent_index, snippet) in snippets {
495 included_parents.insert(parent_index);
496 self.budget_used += snippet.text.len();
497 self.snippets.push(snippet);
498 }
499 }
500
501 fn additional_parent_signatures(
502 &self,
503 path: &Arc<Path>,
504 parent_index: Option<usize>,
505 included_parents: &FxHashSet<usize>,
506 ) -> Result<Vec<(usize, PlannedSnippet<'a>)>> {
507 let mut results = Vec::new();
508 self.additional_parent_signatures_impl(path, parent_index, included_parents, &mut results)?;
509 Ok(results)
510 }
511
512 fn additional_parent_signatures_impl(
513 &self,
514 path: &Arc<Path>,
515 parent_index: Option<usize>,
516 included_parents: &FxHashSet<usize>,
517 results: &mut Vec<(usize, PlannedSnippet<'a>)>,
518 ) -> Result<()> {
519 let Some(parent_index) = parent_index else {
520 return Ok(());
521 };
522 if included_parents.contains(&parent_index) {
523 return Ok(());
524 }
525 let Some(parent_signature) = self.request.signatures.get(parent_index) else {
526 return Err(anyhow!("Invalid parent index {}", parent_index));
527 };
528 results.push((
529 parent_index,
530 PlannedSnippet {
531 path: path.clone(),
532 range: parent_signature.range.clone(),
533 text: &parent_signature.text,
534 text_is_truncated: parent_signature.text_is_truncated,
535 },
536 ));
537 self.additional_parent_signatures_impl(
538 path,
539 parent_signature.parent_index,
540 included_parents,
541 results,
542 )
543 }
544
545 /// Renders the planned context. Each file starts with "```FILE_PATH\n` and ends with triple
546 /// backticks, with a newline after each file. Outputs a line with "..." between nonconsecutive
547 /// chunks.
548 pub fn write(
549 &'a self,
550 excerpt_file_insertions: &mut Vec<(Point, &'static str)>,
551 prompt: &mut String,
552 ) -> Result<SectionLabels> {
553 let mut file_to_snippets: FxHashMap<&'a std::path::Path, Vec<&PlannedSnippet<'a>>> =
554 FxHashMap::default();
555 for snippet in &self.snippets {
556 file_to_snippets
557 .entry(&snippet.path)
558 .or_default()
559 .push(snippet);
560 }
561
562 // Reorder so that file with cursor comes last
563 let mut file_snippets = Vec::new();
564 let mut excerpt_file_snippets = Vec::new();
565 for (file_path, snippets) in file_to_snippets {
566 if file_path == self.request.excerpt_path.as_ref() {
567 excerpt_file_snippets = snippets;
568 } else {
569 file_snippets.push((file_path, snippets, false));
570 }
571 }
572 let excerpt_snippet = PlannedSnippet {
573 path: self.request.excerpt_path.clone(),
574 range: self.request.excerpt_line_range.clone(),
575 text: &self.request.excerpt,
576 text_is_truncated: false,
577 };
578 excerpt_file_snippets.push(&excerpt_snippet);
579 file_snippets.push((&self.request.excerpt_path, excerpt_file_snippets, true));
580
581 let section_labels =
582 self.push_file_snippets(prompt, excerpt_file_insertions, file_snippets)?;
583
584 Ok(section_labels)
585 }
586
587 fn push_file_snippets(
588 &self,
589 output: &mut String,
590 excerpt_file_insertions: &mut Vec<(Point, &'static str)>,
591 file_snippets: Vec<(&'a Path, Vec<&'a PlannedSnippet>, bool)>,
592 ) -> Result<SectionLabels> {
593 let mut section_ranges = Vec::new();
594 let mut excerpt_index = None;
595
596 for (file_path, mut snippets, is_excerpt_file) in file_snippets {
597 snippets.sort_by_key(|s| (s.range.start, Reverse(s.range.end)));
598
599 // TODO: What if the snippets get expanded too large to be editable?
600 let mut current_snippet: Option<(&PlannedSnippet, Range<Line>)> = None;
601 let mut disjoint_snippets: Vec<(&PlannedSnippet, Range<Line>)> = Vec::new();
602 for snippet in snippets {
603 if let Some((_, current_snippet_range)) = current_snippet.as_mut()
604 && snippet.range.start <= current_snippet_range.end
605 {
606 current_snippet_range.end = current_snippet_range.end.max(snippet.range.end);
607 continue;
608 }
609 if let Some(current_snippet) = current_snippet.take() {
610 disjoint_snippets.push(current_snippet);
611 }
612 current_snippet = Some((snippet, snippet.range.clone()));
613 }
614 if let Some(current_snippet) = current_snippet.take() {
615 disjoint_snippets.push(current_snippet);
616 }
617
618 writeln!(output, "`````path={}", file_path.display()).ok();
619 let mut skipped_last_snippet = false;
620 for (snippet, range) in disjoint_snippets {
621 let section_index = section_ranges.len();
622
623 match self.request.prompt_format {
624 PromptFormat::MarkedExcerpt
625 | PromptFormat::OnlySnippets
626 | PromptFormat::NumLinesUniDiff => {
627 if range.start.0 > 0 && !skipped_last_snippet {
628 output.push_str("…\n");
629 }
630 }
631 PromptFormat::LabeledSections => {
632 if is_excerpt_file
633 && range.start <= self.request.excerpt_line_range.start
634 && range.end >= self.request.excerpt_line_range.end
635 {
636 writeln!(output, "<|current_section|>").ok();
637 } else {
638 writeln!(output, "<|section_{}|>", section_index).ok();
639 }
640 }
641 }
642
643 let push_full_snippet = |output: &mut String| {
644 if self.request.prompt_format == PromptFormat::NumLinesUniDiff {
645 for (i, line) in snippet.text.lines().enumerate() {
646 writeln!(output, "{}|{}", i as u32 + range.start.0 + 1, line)?;
647 }
648 } else {
649 output.push_str(&snippet.text);
650 }
651 anyhow::Ok(())
652 };
653
654 if is_excerpt_file {
655 if self.request.prompt_format == PromptFormat::OnlySnippets {
656 if range.start >= self.request.excerpt_line_range.start
657 && range.end <= self.request.excerpt_line_range.end
658 {
659 skipped_last_snippet = true;
660 } else {
661 skipped_last_snippet = false;
662 output.push_str(snippet.text);
663 }
664 } else if !excerpt_file_insertions.is_empty() {
665 let lines = snippet.text.lines().collect::<Vec<_>>();
666 let push_line = |output: &mut String, line_ix: usize| {
667 if self.request.prompt_format == PromptFormat::NumLinesUniDiff {
668 write!(output, "{}|", line_ix as u32 + range.start.0 + 1)?;
669 }
670 anyhow::Ok(writeln!(output, "{}", lines[line_ix])?)
671 };
672 let mut last_line_ix = 0;
673 let mut insertion_ix = 0;
674 while insertion_ix < excerpt_file_insertions.len() {
675 let (point, insertion) = &excerpt_file_insertions[insertion_ix];
676 let found = point.line >= range.start && point.line <= range.end;
677 if found {
678 excerpt_index = Some(section_index);
679 let insertion_line_ix = (point.line.0 - range.start.0) as usize;
680 for line_ix in last_line_ix..insertion_line_ix {
681 push_line(output, line_ix)?;
682 }
683 if let Some(next_line) = lines.get(insertion_line_ix) {
684 if self.request.prompt_format == PromptFormat::NumLinesUniDiff {
685 write!(
686 output,
687 "{}|",
688 insertion_line_ix as u32 + range.start.0 + 1
689 )?
690 }
691 output.push_str(&next_line[..point.column as usize]);
692 output.push_str(insertion);
693 writeln!(output, "{}", &next_line[point.column as usize..])?;
694 } else {
695 writeln!(output, "{}", insertion)?;
696 }
697 last_line_ix = insertion_line_ix + 1;
698 excerpt_file_insertions.remove(insertion_ix);
699 continue;
700 }
701 insertion_ix += 1;
702 }
703 skipped_last_snippet = false;
704 for line_ix in last_line_ix..lines.len() {
705 push_line(output, line_ix)?;
706 }
707 } else {
708 skipped_last_snippet = false;
709 push_full_snippet(output)?;
710 }
711 } else {
712 skipped_last_snippet = false;
713 push_full_snippet(output)?;
714 }
715
716 section_ranges.push((snippet.path.clone(), range));
717 }
718
719 output.push_str("`````\n\n");
720 }
721
722 Ok(SectionLabels {
723 // TODO: Clean this up
724 excerpt_index: match self.request.prompt_format {
725 PromptFormat::OnlySnippets => 0,
726 _ => excerpt_index.context("bug: no snippet found for excerpt")?,
727 },
728 section_ranges,
729 })
730 }
731}
732
733fn declaration_score_density(declaration: &ReferencedDeclaration, style: DeclarationStyle) -> f32 {
734 declaration_score(declaration, style) / declaration_size(declaration, style) as f32
735}
736
737fn declaration_score(declaration: &ReferencedDeclaration, style: DeclarationStyle) -> f32 {
738 match style {
739 DeclarationStyle::Signature => declaration.signature_score,
740 DeclarationStyle::Declaration => declaration.declaration_score,
741 }
742}
743
744fn declaration_size(declaration: &ReferencedDeclaration, style: DeclarationStyle) -> usize {
745 match style {
746 DeclarationStyle::Signature => declaration.signature_range.len(),
747 DeclarationStyle::Declaration => declaration.text.len(),
748 }
749}