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 None,
136 );
137
138 let mut selected = selection.selected;
139
140 crate::postpass::coherence_post_pass(
141 &mut selected,
142 &scoring_result.filtered_fragments,
143 &scoring_result.graph,
144 effective_budget,
145 );
146
147 crate::postpass::rescue_nontrivial_context(
148 &mut selected,
149 &all_fragments,
150 &scoring_result.rel_scores,
151 &core_ids,
152 effective_budget,
153 );
154
155 let used: u32 = selected.iter().map(|f| f.token_count).sum();
156 let remaining = effective_budget.saturating_sub(used);
157 let changed_files: Vec<PathBuf> = changed_paths.iter().map(PathBuf::from).collect();
158 crate::postpass::ensure_changed_files_represented(
159 &mut selected,
160 &all_fragments,
161 &changed_files,
162 remaining,
163 Path::new("."),
164 &[],
165 None,
166 &core_ids,
167 &FxHashMap::default(),
168 );
169
170 let dummy_root = Path::new(".");
171 let mut changed_list: Vec<String> = changed_paths.iter().cloned().collect();
172 changed_list.sort();
173 let change = crate::render::ChangeSummary {
174 lockfile_changes: Vec::new(),
175 commit_message: None,
176 changed_files: changed_list,
177 deleted_files: Vec::new(),
178 renamed_files: Vec::new(),
179 };
180 build_diff_context_output(
181 dummy_root,
182 &selected,
183 no_content,
184 &core_ids,
185 &scoring_result.rel_scores,
186 change,
187 )
188}
189
190fn compute_memory_hunks(
191 initial: &FxHashMap<String, String>,
192 changed: &FxHashMap<String, String>,
193) -> Vec<DiffHunk> {
194 let mut hunks = Vec::new();
195
196 for (path, new_content) in changed {
197 let old_content = initial.get(path).map(|s| s.as_str()).unwrap_or("");
198 if old_content == new_content {
199 continue;
200 }
201 let path_arc: Arc<str> = Arc::from(path.as_str());
202 let file_hunks = diff_to_hunks(&path_arc, old_content, new_content);
203 hunks.extend(file_hunks);
204 }
205
206 for (path, _old_content) in initial {
207 if !changed.contains_key(path) {
208 let path_arc: Arc<str> = Arc::from(path.as_str());
209 let old_line_count = initial[path].lines().count() as u32;
210 if old_line_count > 0 {
211 hunks.push(DiffHunk {
212 path: path_arc,
213 new_start: 1,
214 new_len: 0,
215 old_start: 1,
216 old_len: old_line_count,
217 });
218 }
219 }
220 }
221
222 hunks
223}
224
225fn diff_to_hunks(path: &Arc<str>, old: &str, new: &str) -> Vec<DiffHunk> {
226 let diff = TextDiff::from_lines(old, new);
227 let mut hunks = Vec::new();
228
229 let mut new_line: u32 = 0;
230 let mut old_line: u32 = 0;
231
232 let mut hunk_new_start: Option<u32> = None;
233 let mut hunk_new_len: u32 = 0;
234 let mut hunk_old_start: u32 = 0;
235 let mut hunk_old_len: u32 = 0;
236
237 for change in diff.iter_all_changes() {
238 match change.tag() {
239 ChangeTag::Equal => {
240 if let Some(start) = hunk_new_start.take() {
241 hunks.push(DiffHunk {
242 path: Arc::clone(path),
243 new_start: start,
244 new_len: hunk_new_len,
245 old_start: hunk_old_start,
246 old_len: hunk_old_len,
247 });
248 hunk_new_len = 0;
249 hunk_old_len = 0;
250 }
251 new_line += 1;
252 old_line += 1;
253 }
254 ChangeTag::Delete => {
255 if hunk_new_start.is_none() {
256 hunk_new_start = Some(new_line + 1);
257 hunk_old_start = old_line + 1;
258 }
259 hunk_old_len += 1;
260 old_line += 1;
261 }
262 ChangeTag::Insert => {
263 if hunk_new_start.is_none() {
264 hunk_new_start = Some(new_line + 1);
265 hunk_old_start = old_line + 1;
266 }
267 hunk_new_len += 1;
268 new_line += 1;
269 }
270 }
271 }
272
273 if let Some(start) = hunk_new_start {
274 hunks.push(DiffHunk {
275 path: Arc::clone(path),
276 new_start: start,
277 new_len: hunk_new_len,
278 old_start: hunk_old_start,
279 old_len: hunk_old_len,
280 });
281 }
282
283 hunks
284}
285
286fn compute_memory_diff_text(
287 initial: &FxHashMap<String, String>,
288 changed: &FxHashMap<String, String>,
289) -> String {
290 let mut result = String::new();
291
292 let mut paths: Vec<&String> = changed.keys().collect();
293 paths.sort();
294
295 for path in paths {
296 let new_content = &changed[path];
297 let old_content = initial.get(path).map(|s| s.as_str()).unwrap_or("");
298 if old_content == new_content {
299 continue;
300 }
301
302 let diff = TextDiff::from_lines(old_content, new_content);
303 let mut udiff = diff.unified_diff();
304 let formatted = udiff
305 .context_radius(TOKENIZATION.diff_context_radius)
306 .header(&format!("a/{path}"), &format!("b/{path}"));
307 let _ = write!(result, "{formatted}");
308 }
309
310 let mut deleted_paths: Vec<&String> = initial
311 .keys()
312 .filter(|p| !changed.contains_key(*p))
313 .collect();
314 deleted_paths.sort();
315
316 for path in deleted_paths {
317 let old_content = &initial[path];
318 let empty = String::new();
319 let diff = TextDiff::from_lines(old_content, &empty);
320 let mut udiff = diff.unified_diff();
321 let formatted = udiff
322 .context_radius(TOKENIZATION.diff_context_radius)
323 .header(&format!("a/{path}"), "/dev/null");
324 let _ = write!(result, "{formatted}");
325 }
326
327 result
328}
329
330fn merge_file_contents(
331 initial: &FxHashMap<String, String>,
332 changed: &FxHashMap<String, String>,
333) -> FxHashMap<String, String> {
334 let mut merged = initial.clone();
335 for (path, content) in changed {
336 merged.insert(path.clone(), content.clone());
337 }
338 merged
339}
340
341fn empty_output(name: &str) -> DiffContextOutput {
342 DiffContextOutput {
343 lockfile_changes: Vec::new(),
344 name: name.to_string(),
345 output_type: "diff_context".to_string(),
346 commit_message: None,
347 changed_files: Vec::new(),
348 deleted_files: Vec::new(),
349 renamed_files: Vec::new(),
350 fragment_count: 0,
351 fragments: Vec::new(),
352 latency: None,
353 }
354}