1use std::fmt::Write;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4
5use rayon::prelude::*;
6use rustc_hash::{FxHashMap, FxHashSet};
7use similar::{ChangeTag, TextDiff};
8
9use crate::config::budget::BUDGET;
10use crate::config::limits::LIMITS;
11use crate::config::tokenization::TOKENIZATION;
12use crate::core::{compute_seed_weights, identify_core_fragments};
13use crate::edges;
14use crate::mode::{PipelineConfig, ScoringKind, ScoringMode};
15use crate::parsers::fragment_file;
16use crate::render::{DiffContextOutput, build_diff_context_output};
17use crate::scoring::{BM25Scoring, EgoGraphScoring, PPRScoring, ScoringStrategy};
18use crate::signatures::generate_signature_variants;
19use crate::tokenizer::count_tokens;
20use crate::types::{DiffHunk, Fragment, FragmentId};
21
22pub struct MemoryRepo {
23 pub name: String,
24 pub initial_files: FxHashMap<String, String>,
25 pub changed_files: FxHashMap<String, String>,
26}
27
28pub fn build_diff_context_in_memory(
29 repo: &MemoryRepo,
30 budget_tokens: Option<u32>,
31 _alpha: f64,
32 tau: f64,
33 no_content: bool,
34 scoring_mode: ScoringMode,
35) -> DiffContextOutput {
36 let hunks = compute_memory_hunks(&repo.initial_files, &repo.changed_files);
37 if hunks.is_empty() {
38 return empty_output(&repo.name);
39 }
40
41 let diff_text = compute_memory_diff_text(&repo.initial_files, &repo.changed_files);
42 let all_files = merge_file_contents(&repo.initial_files, &repo.changed_files);
43
44 let changed_paths: FxHashSet<String> =
45 hunks.iter().map(|h| h.path.as_ref().to_string()).collect();
46
47 let changed_file_paths: Vec<PathBuf> = changed_paths.iter().map(PathBuf::from).collect();
48 let all_file_paths: Vec<PathBuf> = all_files.keys().map(PathBuf::from).collect();
49 let file_cache: FxHashMap<PathBuf, String> = all_files
50 .iter()
51 .map(|(k, v)| (PathBuf::from(k), v.clone()))
52 .collect();
53
54 let discovered = edges::discover_all_related_files(
55 &changed_file_paths,
56 &all_file_paths,
57 None,
58 Some(&file_cache),
59 );
60 let discovered_paths: FxHashSet<String> = discovered
61 .iter()
62 .map(|p| p.to_string_lossy().to_string())
63 .collect();
64
65 let allowed_paths: FxHashSet<&str> = changed_paths
66 .iter()
67 .chain(discovered_paths.iter())
68 .map(|s| s.as_str())
69 .collect();
70
71 let mut all_fragments: Vec<Fragment> = Vec::new();
72 let mut seen: FxHashSet<FragmentId> = FxHashSet::default();
73 for (path, content) in &all_files {
74 if !allowed_paths.contains(path.as_str()) {
75 continue;
76 }
77 let path_arc: Arc<str> = Arc::from(path.as_str());
78 let frags = fragment_file(path_arc, content);
79 for f in frags {
80 if seen.insert(f.id.clone()) {
81 all_fragments.push(f);
82 }
83 }
84 }
85
86 all_fragments.par_iter_mut().for_each(|f| {
87 f.token_count = count_tokens(&f.content) + LIMITS.overhead_per_fragment;
88 });
89
90 let core_ids = identify_core_fragments(&hunks, &all_fragments);
91
92 let mut sig_frags = generate_signature_variants(&all_fragments);
93 sig_frags.par_iter_mut().for_each(|f| {
94 f.token_count = count_tokens(&f.content) + LIMITS.overhead_per_fragment;
95 });
96 all_fragments.extend(sig_frags);
97
98 let effective_budget = budget_tokens.unwrap_or(BUDGET.unlimited);
99 let config = PipelineConfig::from_mode(scoring_mode);
100 let seed_weights = compute_seed_weights(&hunks, &core_ids, &all_fragments);
101
102 let discovered_arc: FxHashSet<Arc<str>> = discovered_paths
103 .iter()
104 .map(|s| Arc::from(s.as_str()))
105 .collect();
106
107 let strategy: Box<dyn ScoringStrategy> = match config.scoring {
108 ScoringKind::Ego => Box::new(EgoGraphScoring::new(config.ego_depth)),
109 ScoringKind::Ppr => Box::new(PPRScoring::new(
110 config.ppr_alpha,
111 config.low_relevance_filter,
112 )),
113 ScoringKind::Bm25 => Box::new(BM25Scoring),
114 };
115
116 let scoring_result = strategy.score_and_filter(
117 &all_fragments,
118 &core_ids,
119 &hunks,
120 None,
121 Some(&seed_weights),
122 Some(&discovered_arc),
123 );
124
125 let needs = crate::utility::needs::needs_from_diff(&all_fragments, &core_ids, &diff_text);
126
127 let selection = crate::select::lazy_greedy_select(
128 scoring_result.filtered_fragments.clone(),
129 &core_ids,
130 &scoring_result.rel_scores,
131 &needs,
132 effective_budget,
133 tau,
134 None,
135 );
136
137 let mut selected = selection.selected;
138
139 crate::postpass::coherence_post_pass(
140 &mut selected,
141 &scoring_result.filtered_fragments,
142 &scoring_result.graph,
143 effective_budget,
144 );
145
146 crate::postpass::rescue_nontrivial_context(
147 &mut selected,
148 &all_fragments,
149 &scoring_result.rel_scores,
150 &core_ids,
151 effective_budget,
152 );
153
154 let used: u32 = selected.iter().map(|f| f.token_count).sum();
155 let remaining = effective_budget.saturating_sub(used);
156 let changed_files: Vec<PathBuf> = changed_paths.iter().map(PathBuf::from).collect();
157 crate::postpass::ensure_changed_files_represented(
158 &mut selected,
159 &all_fragments,
160 &changed_files,
161 remaining,
162 Path::new("."),
163 &[],
164 None,
165 &core_ids,
166 );
167
168 let dummy_root = Path::new(".");
169 let mut changed_list: Vec<String> = changed_paths.iter().cloned().collect();
170 changed_list.sort();
171 let change = crate::render::ChangeSummary {
172 commit_message: None,
173 changed_files: changed_list,
174 deleted_files: Vec::new(),
175 renamed_files: Vec::new(),
176 };
177 build_diff_context_output(
178 dummy_root,
179 &selected,
180 no_content,
181 &core_ids,
182 &scoring_result.rel_scores,
183 change,
184 )
185}
186
187fn compute_memory_hunks(
188 initial: &FxHashMap<String, String>,
189 changed: &FxHashMap<String, String>,
190) -> Vec<DiffHunk> {
191 let mut hunks = Vec::new();
192
193 for (path, new_content) in changed {
194 let old_content = initial.get(path).map(|s| s.as_str()).unwrap_or("");
195 if old_content == new_content {
196 continue;
197 }
198 let path_arc: Arc<str> = Arc::from(path.as_str());
199 let file_hunks = diff_to_hunks(&path_arc, old_content, new_content);
200 hunks.extend(file_hunks);
201 }
202
203 for (path, _old_content) in initial {
204 if !changed.contains_key(path) {
205 let path_arc: Arc<str> = Arc::from(path.as_str());
206 let old_line_count = initial[path].lines().count() as u32;
207 if old_line_count > 0 {
208 hunks.push(DiffHunk {
209 path: path_arc,
210 new_start: 1,
211 new_len: 0,
212 old_start: 1,
213 old_len: old_line_count,
214 });
215 }
216 }
217 }
218
219 hunks
220}
221
222fn diff_to_hunks(path: &Arc<str>, old: &str, new: &str) -> Vec<DiffHunk> {
223 let diff = TextDiff::from_lines(old, new);
224 let mut hunks = Vec::new();
225
226 let mut new_line: u32 = 0;
227 let mut old_line: u32 = 0;
228
229 let mut hunk_new_start: Option<u32> = None;
230 let mut hunk_new_len: u32 = 0;
231 let mut hunk_old_start: u32 = 0;
232 let mut hunk_old_len: u32 = 0;
233
234 for change in diff.iter_all_changes() {
235 match change.tag() {
236 ChangeTag::Equal => {
237 if let Some(start) = hunk_new_start.take() {
238 hunks.push(DiffHunk {
239 path: Arc::clone(path),
240 new_start: start,
241 new_len: hunk_new_len,
242 old_start: hunk_old_start,
243 old_len: hunk_old_len,
244 });
245 hunk_new_len = 0;
246 hunk_old_len = 0;
247 }
248 new_line += 1;
249 old_line += 1;
250 }
251 ChangeTag::Delete => {
252 if hunk_new_start.is_none() {
253 hunk_new_start = Some(new_line + 1);
254 hunk_old_start = old_line + 1;
255 }
256 hunk_old_len += 1;
257 old_line += 1;
258 }
259 ChangeTag::Insert => {
260 if hunk_new_start.is_none() {
261 hunk_new_start = Some(new_line + 1);
262 hunk_old_start = old_line + 1;
263 }
264 hunk_new_len += 1;
265 new_line += 1;
266 }
267 }
268 }
269
270 if let Some(start) = hunk_new_start {
271 hunks.push(DiffHunk {
272 path: Arc::clone(path),
273 new_start: start,
274 new_len: hunk_new_len,
275 old_start: hunk_old_start,
276 old_len: hunk_old_len,
277 });
278 }
279
280 hunks
281}
282
283fn compute_memory_diff_text(
284 initial: &FxHashMap<String, String>,
285 changed: &FxHashMap<String, String>,
286) -> String {
287 let mut result = String::new();
288
289 let mut paths: Vec<&String> = changed.keys().collect();
290 paths.sort();
291
292 for path in paths {
293 let new_content = &changed[path];
294 let old_content = initial.get(path).map(|s| s.as_str()).unwrap_or("");
295 if old_content == new_content {
296 continue;
297 }
298
299 let diff = TextDiff::from_lines(old_content, new_content);
300 let mut udiff = diff.unified_diff();
301 let formatted = udiff
302 .context_radius(TOKENIZATION.diff_context_radius)
303 .header(&format!("a/{path}"), &format!("b/{path}"));
304 let _ = write!(result, "{formatted}");
305 }
306
307 let mut deleted_paths: Vec<&String> = initial
308 .keys()
309 .filter(|p| !changed.contains_key(*p))
310 .collect();
311 deleted_paths.sort();
312
313 for path in deleted_paths {
314 let old_content = &initial[path];
315 let empty = String::new();
316 let diff = TextDiff::from_lines(old_content, &empty);
317 let mut udiff = diff.unified_diff();
318 let formatted = udiff
319 .context_radius(TOKENIZATION.diff_context_radius)
320 .header(&format!("a/{path}"), "/dev/null");
321 let _ = write!(result, "{formatted}");
322 }
323
324 result
325}
326
327fn merge_file_contents(
328 initial: &FxHashMap<String, String>,
329 changed: &FxHashMap<String, String>,
330) -> FxHashMap<String, String> {
331 let mut merged = initial.clone();
332 for (path, content) in changed {
333 merged.insert(path.clone(), content.clone());
334 }
335 merged
336}
337
338fn empty_output(name: &str) -> DiffContextOutput {
339 DiffContextOutput {
340 name: name.to_string(),
341 output_type: "diff_context".to_string(),
342 commit_message: None,
343 changed_files: Vec::new(),
344 deleted_files: Vec::new(),
345 renamed_files: Vec::new(),
346 fragment_count: 0,
347 fragments: Vec::new(),
348 latency: None,
349 }
350}