1use cloud_llm_client::predict_edits_v3::ScoreComponents;
2use itertools::Itertools as _;
3use language::BufferSnapshot;
4use ordered_float::OrderedFloat;
5use serde::Serialize;
6use std::{collections::HashMap, ops::Range};
7use strum::EnumIter;
8use text::{Point, ToPoint};
9
10use crate::{
11 Declaration, EditPredictionExcerpt, EditPredictionExcerptText, Identifier,
12 reference::{Reference, ReferenceRegion},
13 syntax_index::SyntaxIndexState,
14 text_similarity::{IdentifierOccurrences, jaccard_similarity, weighted_overlap_coefficient},
15};
16
17const MAX_IDENTIFIER_DECLARATION_COUNT: usize = 16;
18
19#[derive(Clone, Debug)]
20pub struct ScoredSnippet {
21 pub identifier: Identifier,
22 pub declaration: Declaration,
23 pub score_components: ScoreComponents,
24 pub scores: Scores,
25}
26
27#[derive(EnumIter, Clone, Copy, PartialEq, Eq, Hash, Debug)]
28pub enum SnippetStyle {
29 Signature,
30 Declaration,
31}
32
33impl ScoredSnippet {
34 /// Returns the score for this snippet with the specified style.
35 pub fn score(&self, style: SnippetStyle) -> f32 {
36 match style {
37 SnippetStyle::Signature => self.scores.signature,
38 SnippetStyle::Declaration => self.scores.declaration,
39 }
40 }
41
42 pub fn size(&self, style: SnippetStyle) -> usize {
43 match &self.declaration {
44 Declaration::File { declaration, .. } => match style {
45 SnippetStyle::Signature => declaration.signature_range.len(),
46 SnippetStyle::Declaration => declaration.text.len(),
47 },
48 Declaration::Buffer { declaration, .. } => match style {
49 SnippetStyle::Signature => declaration.signature_range.len(),
50 SnippetStyle::Declaration => declaration.item_range.len(),
51 },
52 }
53 }
54
55 pub fn score_density(&self, style: SnippetStyle) -> f32 {
56 self.score(style) / (self.size(style)) as f32
57 }
58}
59
60pub fn scored_snippets(
61 index: &SyntaxIndexState,
62 excerpt: &EditPredictionExcerpt,
63 excerpt_text: &EditPredictionExcerptText,
64 identifier_to_references: HashMap<Identifier, Vec<Reference>>,
65 cursor_offset: usize,
66 current_buffer: &BufferSnapshot,
67) -> Vec<ScoredSnippet> {
68 let containing_range_identifier_occurrences =
69 IdentifierOccurrences::within_string(&excerpt_text.body);
70 let cursor_point = cursor_offset.to_point(¤t_buffer);
71
72 let start_point = Point::new(cursor_point.row.saturating_sub(2), 0);
73 let end_point = Point::new(cursor_point.row + 1, 0);
74 let adjacent_identifier_occurrences = IdentifierOccurrences::within_string(
75 ¤t_buffer
76 .text_for_range(start_point..end_point)
77 .collect::<String>(),
78 );
79
80 let mut snippets = identifier_to_references
81 .into_iter()
82 .flat_map(|(identifier, references)| {
83 let declarations =
84 index.declarations_for_identifier::<MAX_IDENTIFIER_DECLARATION_COUNT>(&identifier);
85 let declaration_count = declarations.len();
86
87 declarations
88 .into_iter()
89 .filter_map(|(declaration_id, declaration)| match declaration {
90 Declaration::Buffer {
91 buffer_id,
92 declaration: buffer_declaration,
93 ..
94 } => {
95 let is_same_file = buffer_id == ¤t_buffer.remote_id();
96
97 if is_same_file {
98 let overlaps_excerpt =
99 range_intersection(&buffer_declaration.item_range, &excerpt.range)
100 .is_some();
101 if overlaps_excerpt
102 || excerpt
103 .parent_declarations
104 .iter()
105 .any(|(excerpt_parent, _)| excerpt_parent == &declaration_id)
106 {
107 None
108 } else {
109 let declaration_line = buffer_declaration
110 .item_range
111 .start
112 .to_point(current_buffer)
113 .row;
114 Some((
115 true,
116 (cursor_point.row as i32 - declaration_line as i32)
117 .unsigned_abs(),
118 declaration,
119 ))
120 }
121 } else {
122 Some((false, u32::MAX, declaration))
123 }
124 }
125 Declaration::File { .. } => {
126 // We can assume that a file declaration is in a different file,
127 // because the current one must be open
128 Some((false, u32::MAX, declaration))
129 }
130 })
131 .sorted_by_key(|&(_, distance, _)| distance)
132 .enumerate()
133 .map(
134 |(
135 declaration_line_distance_rank,
136 (is_same_file, declaration_line_distance, declaration),
137 )| {
138 let same_file_declaration_count = index.file_declaration_count(declaration);
139
140 score_snippet(
141 &identifier,
142 &references,
143 declaration.clone(),
144 is_same_file,
145 declaration_line_distance,
146 declaration_line_distance_rank,
147 same_file_declaration_count,
148 declaration_count,
149 &containing_range_identifier_occurrences,
150 &adjacent_identifier_occurrences,
151 cursor_point,
152 current_buffer,
153 )
154 },
155 )
156 .collect::<Vec<_>>()
157 })
158 .flatten()
159 .collect::<Vec<_>>();
160
161 snippets.sort_unstable_by_key(|snippet| {
162 OrderedFloat(
163 snippet
164 .score_density(SnippetStyle::Declaration)
165 .max(snippet.score_density(SnippetStyle::Signature)),
166 )
167 });
168
169 snippets
170}
171
172fn range_intersection<T: Ord + Clone>(a: &Range<T>, b: &Range<T>) -> Option<Range<T>> {
173 let start = a.start.clone().max(b.start.clone());
174 let end = a.end.clone().min(b.end.clone());
175 if start < end {
176 Some(Range { start, end })
177 } else {
178 None
179 }
180}
181
182fn score_snippet(
183 identifier: &Identifier,
184 references: &[Reference],
185 declaration: Declaration,
186 is_same_file: bool,
187 declaration_line_distance: u32,
188 declaration_line_distance_rank: usize,
189 same_file_declaration_count: usize,
190 declaration_count: usize,
191 containing_range_identifier_occurrences: &IdentifierOccurrences,
192 adjacent_identifier_occurrences: &IdentifierOccurrences,
193 cursor: Point,
194 current_buffer: &BufferSnapshot,
195) -> Option<ScoredSnippet> {
196 let is_referenced_nearby = references
197 .iter()
198 .any(|r| r.region == ReferenceRegion::Nearby);
199 let is_referenced_in_breadcrumb = references
200 .iter()
201 .any(|r| r.region == ReferenceRegion::Breadcrumb);
202 let reference_count = references.len();
203 let reference_line_distance = references
204 .iter()
205 .map(|r| {
206 let reference_line = r.range.start.to_point(current_buffer).row as i32;
207 (cursor.row as i32 - reference_line).unsigned_abs()
208 })
209 .min()
210 .unwrap();
211
212 let item_source_occurrences = IdentifierOccurrences::within_string(&declaration.item_text().0);
213 let item_signature_occurrences =
214 IdentifierOccurrences::within_string(&declaration.signature_text().0);
215 let containing_range_vs_item_jaccard = jaccard_similarity(
216 containing_range_identifier_occurrences,
217 &item_source_occurrences,
218 );
219 let containing_range_vs_signature_jaccard = jaccard_similarity(
220 containing_range_identifier_occurrences,
221 &item_signature_occurrences,
222 );
223 let adjacent_vs_item_jaccard =
224 jaccard_similarity(adjacent_identifier_occurrences, &item_source_occurrences);
225 let adjacent_vs_signature_jaccard =
226 jaccard_similarity(adjacent_identifier_occurrences, &item_signature_occurrences);
227
228 let containing_range_vs_item_weighted_overlap = weighted_overlap_coefficient(
229 containing_range_identifier_occurrences,
230 &item_source_occurrences,
231 );
232 let containing_range_vs_signature_weighted_overlap = weighted_overlap_coefficient(
233 containing_range_identifier_occurrences,
234 &item_signature_occurrences,
235 );
236 let adjacent_vs_item_weighted_overlap =
237 weighted_overlap_coefficient(adjacent_identifier_occurrences, &item_source_occurrences);
238 let adjacent_vs_signature_weighted_overlap =
239 weighted_overlap_coefficient(adjacent_identifier_occurrences, &item_signature_occurrences);
240
241 // TODO: Consider adding declaration_file_count
242 let score_components = ScoreComponents {
243 is_same_file,
244 is_referenced_nearby,
245 is_referenced_in_breadcrumb,
246 reference_line_distance,
247 declaration_line_distance,
248 declaration_line_distance_rank,
249 reference_count,
250 same_file_declaration_count,
251 declaration_count,
252 containing_range_vs_item_jaccard,
253 containing_range_vs_signature_jaccard,
254 adjacent_vs_item_jaccard,
255 adjacent_vs_signature_jaccard,
256 containing_range_vs_item_weighted_overlap,
257 containing_range_vs_signature_weighted_overlap,
258 adjacent_vs_item_weighted_overlap,
259 adjacent_vs_signature_weighted_overlap,
260 };
261
262 Some(ScoredSnippet {
263 identifier: identifier.clone(),
264 declaration: declaration,
265 scores: Scores::score(&score_components),
266 score_components,
267 })
268}
269
270#[derive(Clone, Debug, Serialize)]
271pub struct Scores {
272 pub signature: f32,
273 pub declaration: f32,
274}
275
276impl Scores {
277 fn score(components: &ScoreComponents) -> Scores {
278 // TODO: handle truncation
279
280 // Score related to how likely this is the correct declaration, range 0 to 1
281 let accuracy_score = if components.is_same_file {
282 // TODO: use declaration_line_distance_rank
283 1.0 / components.same_file_declaration_count as f32
284 } else {
285 1.0 / components.declaration_count as f32
286 };
287
288 // Score related to the distance between the reference and cursor, range 0 to 1
289 let distance_score = if components.is_referenced_nearby {
290 1.0 / (1.0 + components.reference_line_distance as f32 / 10.0).powf(2.0)
291 } else {
292 // same score as ~14 lines away, rationale is to not overly penalize references from parent signatures
293 0.5
294 };
295
296 // For now instead of linear combination, the scores are just multiplied together.
297 let combined_score = 10.0 * accuracy_score * distance_score;
298
299 Scores {
300 signature: combined_score * components.containing_range_vs_signature_weighted_overlap,
301 // declaration score gets boosted both by being multiplied by 2 and by there being more
302 // weighted overlap.
303 declaration: 2.0
304 * combined_score
305 * components.containing_range_vs_item_weighted_overlap,
306 }
307 }
308}