1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3use std::time::Instant;
4
5use anyhow::Result;
6use rayon::prelude::*;
7use rustc_hash::{FxHashMap, FxHashSet};
8
9use crate::candidate_files;
10use crate::config::budget::BUDGET;
11use crate::config::graph_filtering::GRAPH_FILTERING;
12use crate::config::limits::LIMITS;
13use crate::config::tokenization::TOKENIZATION;
14use crate::core::{compute_seed_weights, identify_core_fragments};
15use crate::discovery::{
16 BM25Discovery, DefaultDiscovery, DiscoveryContext, DiscoveryStrategy, EnsembleDiscovery,
17 TestFileDiscovery,
18};
19use crate::fragmentation::process_files_for_fragments;
20use crate::git::{self, CatFileBatch};
21use crate::mode::{DiscoveryKind, PipelineConfig, ScoringMode};
22use crate::postpass;
23use crate::render::{self, DiffContextOutput};
24use crate::scoring::{ScoringResult, create_scoring_strategy};
25use crate::signatures::generate_signature_variants;
26use crate::tokenizer::count_tokens;
27use crate::types::{Fragment, FragmentId};
28use crate::utility::InformationNeed;
29
30pub struct ScoredState {
37 pub root_dir: PathBuf,
38 pub config: PipelineConfig,
39 pub all_fragments: Vec<Fragment>,
40 pub core_ids: FxHashSet<FragmentId>,
41 pub core_excerpts: FxHashMap<FragmentId, Fragment>,
45 pub scoring_result: ScoringResult,
46 pub needs: Vec<InformationNeed>,
47 pub changed_files: Vec<PathBuf>,
48 pub deleted_files: Vec<String>,
49 pub renamed_files: Vec<(String, String)>,
50 pub lockfile_changes: Vec<String>,
53 pub ignored_changes: Vec<String>,
59 pub policy_excluded_count: usize,
63 pub discovery_source: FxHashMap<Arc<str>, &'static str>,
70 pub preferred_revs: Vec<String>,
71 pub commit_message: Option<String>,
72 pub heavy_latency_ms: HeavyLatencyMs,
73}
74
75#[derive(Default, Clone, Copy)]
76pub struct HeavyLatencyMs {
77 pub pre_phase: f64,
82 pub parse_changed: f64,
83 pub universe_walk: f64,
84 pub discovery: f64,
85 pub parse_discovered: f64,
86 pub tokenization: f64,
87 pub graph_build: f64,
88 pub scoring: f64,
89}
90
91pub fn build_diff_context(
92 root_dir: &Path,
93 diff_range: Option<&str>,
94 budget_tokens: Option<u32>,
95 alpha: f64,
96 tau: f64,
97 no_content: bool,
98 full: bool,
99 scoring_mode: ScoringMode,
100 timeout: u64,
101) -> Result<DiffContextOutput> {
102 if full {
103 return build_diff_context_full(root_dir, diff_range, no_content, timeout);
104 }
105 let state = compute_scored_state(root_dir, diff_range, alpha, scoring_mode, timeout)?;
106 if state.all_fragments.is_empty() {
107 return Ok(empty_output_from_state(&state));
108 }
109 Ok(select_with_params(&state, budget_tokens, tau, no_content))
110}
111
112pub fn build_diff_context_locate(
116 root_dir: &Path,
117 diff_range: Option<&str>,
118 budget_tokens: Option<u32>,
119 alpha: f64,
120 tau: f64,
121 scoring_mode: ScoringMode,
122 timeout: u64,
123) -> Result<crate::locate::LocateOutput> {
124 let state = compute_scored_state(root_dir, diff_range, alpha, scoring_mode, timeout)?;
125 let outcome = if state.all_fragments.is_empty() {
126 SelectionOutcome {
127 selected: Vec::new(),
128 effective_budget: budget_tokens.unwrap_or(0),
129 selection_iters: 0,
130 stopping_certificate: 0.0,
131 select_ms: 0.0,
132 }
133 } else {
134 run_selection(&state, budget_tokens, tau)
135 };
136 Ok(crate::locate::build_locate(&state, &outcome))
137}
138
139fn count_text_lines(path: &Path) -> Option<u32> {
161 use std::io::Read;
162
163 const CHUNK: usize = 64 * 1024;
164
165 let mut file = std::fs::File::open(path).ok()?;
166 let mut buf = vec![0u8; CHUNK];
167 let mut carry: Vec<u8> = Vec::new();
168 let mut newlines: u32 = 0;
169 let mut last_byte: Option<u8> = None;
170
171 loop {
172 let read = file.read(&mut buf).ok()?;
173 if read == 0 {
174 break;
175 }
176 carry.extend_from_slice(&buf[..read]);
177 let valid_upto = match std::str::from_utf8(&carry) {
178 Ok(_) => carry.len(),
179 Err(e) if e.error_len().is_none() => e.valid_up_to(),
182 Err(_) => return None,
183 };
184 newlines = newlines
185 .saturating_add(carry[..valid_upto].iter().filter(|b| **b == b'\n').count() as u32);
186 if let Some(&b) = carry[..valid_upto].last() {
187 last_byte = Some(b);
188 }
189 carry.drain(..valid_upto);
190 }
191
192 if !carry.is_empty() {
195 return None;
196 }
197 Some(match last_byte {
200 None => 0,
201 Some(b'\n') => newlines,
202 Some(_) => newlines.saturating_add(1),
203 })
204}
205
206pub fn compute_scored_state(
210 root_dir: &Path,
211 diff_range: Option<&str>,
212 alpha: f64,
213 scoring_mode: ScoringMode,
214 timeout: u64,
215) -> Result<ScoredState> {
216 let t_entry = Instant::now();
217 git::set_git_timeout(timeout);
218 crate::deadline::set_compute_deadline(timeout);
219 let root_dir = resolve_repo_root(root_dir)?;
220 if alpha <= 0.0 || alpha >= 1.0 {
221 anyhow::bail!("alpha must be in (0, 1), got {}", alpha);
222 }
223
224 let resolved = git::resolve_duration_range(&root_dir, diff_range)?;
225 let diff_range = resolved.range.as_deref();
226
227 let mut hunks = git::parse_diff(&root_dir, diff_range)?;
228
229 let is_working_tree_diff = resolved.from_duration || matches!(diff_range, None | Some("HEAD"));
235 let mut untracked_files: Vec<PathBuf> = Vec::new();
236 if is_working_tree_diff {
237 if let Ok(files) = git::get_untracked_files(&root_dir) {
238 for f in &files {
239 if let Some(line_count) = count_text_lines(f) {
240 if line_count > 0 {
241 let path_str: Arc<str> = Arc::from(f.to_string_lossy().as_ref());
242 hunks.push(crate::types::DiffHunk {
243 path: path_str,
244 new_start: 1,
245 new_len: line_count,
246 old_start: 0,
247 old_len: 0,
248 });
249 }
250 }
251 }
252 untracked_files = files;
253 }
254 }
255
256 hunks.retain(|h| !is_secret_path(Path::new(&*h.path)));
257
258 if hunks.is_empty() {
259 return Ok(empty_scored_state_with_changes(root_dir, diff_range));
260 }
261
262 let ignored_rel_paths = resolve_ignored_paths(&root_dir, &hunks);
263 let mut ignored_display: Vec<String> = Vec::new();
269 let mut policy_excluded_paths: FxHashSet<String> = FxHashSet::default();
272 for h in &hunks {
273 let p = Path::new(&*h.path);
274 let Some(rel) = rel_path_string(&root_dir, p) else {
275 continue;
276 };
277 match ignored_rel_paths.get(&rel) {
278 Some(git::IgnoreSource::Gitignore) => ignored_display.push(rel),
279 Some(git::IgnoreSource::DiffctxPolicy) => {
280 policy_excluded_paths.insert(rel);
281 }
282 None => {}
283 }
284 }
285 let policy_excluded = policy_excluded_paths.len();
286 ignored_display.sort();
287 ignored_display.dedup();
288 hunks.retain(|h| !is_ignored_path(&root_dir, Path::new(&*h.path), &ignored_rel_paths));
289
290 let mut lockfile_display: Vec<String> = hunks
291 .iter()
292 .filter(|h| is_lockfile_path(Path::new(&*h.path)))
293 .filter_map(|h| rel_path_string(&root_dir, Path::new(&*h.path)))
294 .collect();
295 lockfile_display.sort();
296 lockfile_display.dedup();
297 hunks.retain(|h| !is_lockfile_path(Path::new(&*h.path)));
298
299 if hunks.is_empty() {
300 let mut state = empty_scored_state_with_changes(root_dir, diff_range);
301 state.lockfile_changes = lockfile_display;
302 state.ignored_changes = ignored_display;
303 state.policy_excluded_count = policy_excluded;
304 return Ok(state);
305 }
306
307 let diff_text = git::get_diff_text(&root_dir, diff_range)?;
308
309 let mut changed_files = git::get_changed_files(&root_dir, diff_range)?;
310 changed_files.extend(untracked_files);
311 if changed_files.is_empty() {
312 return Ok(empty_scored_state_with_changes(root_dir, diff_range));
313 }
314
315 let deleted_files = git::get_deleted_files(&root_dir, diff_range)?;
316 let renamed_old = git::get_renamed_paths(&root_dir, diff_range)?;
320 let mut deleted_display: Vec<String> = deleted_files
324 .iter()
325 .map(|p| {
326 p.strip_prefix(&root_dir)
327 .unwrap_or(p)
328 .to_string_lossy()
329 .replace('\\', "/")
330 })
331 .collect();
332 deleted_display.sort();
333 let renamed_display = git::get_rename_pairs(&root_dir, diff_range).unwrap_or_default();
334 let excluded: FxHashSet<PathBuf> = deleted_files.into_iter().chain(renamed_old).collect();
335 let changed_files: Vec<PathBuf> = changed_files
336 .into_iter()
337 .filter(|f| {
338 let resolved = f.canonicalize().unwrap_or_else(|_| f.clone());
339 !excluded.contains(&resolved)
340 && !is_secret_path(f)
341 && !is_lockfile_path(f)
342 && !is_ignored_path(&root_dir, f, &ignored_rel_paths)
343 })
344 .collect();
345
346 let (base_rev, head_rev) = diff_range
347 .map(git::split_diff_range)
348 .unwrap_or((None, None));
349 let preferred_revs = build_preferred_revs(base_rev.as_deref(), head_rev.as_deref());
350 let commit_message = head_rev
351 .as_deref()
352 .and_then(|h| git::get_commit_message(&root_dir, h).ok())
353 .and_then(|m| {
354 m.lines()
355 .map(str::trim)
356 .find(|l| !l.is_empty())
357 .map(str::to_string)
358 });
359
360 let t0 = Instant::now();
361 let pre_phase_ms = t0.duration_since(t_entry).as_secs_f64() * 1000.0;
362
363 let mut seen_frag_ids: FxHashSet<FragmentId> = FxHashSet::default();
364 let mut batch_reader = CatFileBatch::new(&root_dir)?;
365 let mut all_fragments = process_files_for_fragments(
366 &changed_files,
367 &root_dir,
368 &preferred_revs,
369 &mut seen_frag_ids,
370 Some(&mut batch_reader),
371 true,
372 );
373
374 let t_parse_changed = Instant::now();
375
376 let included_set: FxHashSet<PathBuf> = changed_files.iter().cloned().collect();
377 let all_candidate_files = candidate_files::collect_candidate_files(&root_dir, &included_set);
378
379 let t_universe = Instant::now();
380
381 let file_cache = build_file_cache(&all_candidate_files);
382 let mode = scoring_mode;
383 let mut config = PipelineConfig::from_mode(mode);
384 if let Ok(s) = std::env::var("DIFFCTX_OBJECTIVE") {
385 config.objective = crate::mode::ObjectiveMode::from_str(&s);
386 }
387
388 let mut expansion_concepts: FxHashSet<String> =
389 crate::types::extract_identifiers(&diff_text, TOKENIZATION.query_min_identifier_length)
390 .into_iter()
391 .collect();
392
393 if let Some(ref h) = head_rev {
394 if std::env::var("DIFFCTX_NO_COMMIT_SIGNAL").as_deref() != Ok("1") {
395 if let Ok(commit_msg) = git::get_commit_message(&root_dir, h) {
396 for ident in crate::types::extract_identifiers(
397 &commit_msg,
398 TOKENIZATION.query_min_identifier_length,
399 ) {
400 expansion_concepts.insert(ident);
401 }
402 }
403 }
404 }
405
406 let discovery_ctx = DiscoveryContext {
407 root_dir: root_dir.clone(),
408 changed_files: changed_files.clone(),
409 all_candidates: all_candidate_files,
410 diff_text: diff_text.clone(),
411 expansion_concepts,
412 file_cache,
413 token_corpus: std::sync::OnceLock::new(),
414 };
415
416 let (discovered_files, discovery_attribution) =
417 create_discovery(&config).discover_attributed(&discovery_ctx);
418 let discovered_files: Vec<PathBuf> = discovered_files
419 .into_iter()
420 .map(|p| candidate_files::normalize_path(&p, &root_dir))
421 .collect();
422 let discovery_source: FxHashMap<Arc<str>, &'static str> = discovery_attribution
423 .into_iter()
424 .map(|(path, source)| {
425 let normalized = candidate_files::normalize_path(&path, &root_dir);
426 (Arc::from(normalized.to_string_lossy().as_ref()), source)
427 })
428 .collect();
429
430 drop(discovery_ctx);
431
432 let t_discovery = Instant::now();
433
434 all_fragments.extend(process_files_for_fragments(
435 &discovered_files,
436 &root_dir,
437 &preferred_revs,
438 &mut seen_frag_ids,
439 Some(&mut batch_reader),
440 false,
441 ));
442
443 let t_parse_discovered = Instant::now();
444
445 assign_token_counts(&mut all_fragments);
446
447 let core_ids = identify_core_fragments(&hunks, &all_fragments);
448
449 let mut core_excerpts =
450 crate::excerpt::generate_core_excerpts(&all_fragments, &core_ids, &hunks);
451 assign_excerpt_token_counts(&mut core_excerpts);
452
453 let signature_frags = generate_signature_variants(&all_fragments);
454 let mut sig_frags = signature_frags;
455 assign_token_counts(&mut sig_frags);
456 all_fragments.extend(sig_frags);
457
458 let t_tokenization = Instant::now();
459
460 let seed_weights = compute_seed_weights(&hunks, &core_ids, &all_fragments);
461
462 let discovered_path_set: FxHashSet<Arc<str>> = discovered_files
463 .iter()
464 .map(|p| Arc::from(p.to_string_lossy().as_ref()))
465 .collect();
466
467 let strategy = create_scoring_strategy(&config);
468
469 let scoring_result = strategy.score_and_filter(
470 &all_fragments,
471 &core_ids,
472 &hunks,
473 Some(root_dir.as_path()),
474 Some(&seed_weights),
475 Some(&discovered_path_set),
476 );
477
478 let needs = crate::utility::needs::needs_from_diff(&all_fragments, &core_ids, &diff_text);
479
480 let t_done = Instant::now();
481 batch_reader.close();
482
483 let graph_build_ms = scoring_result.graph_build_ms;
484 let heavy_latency_ms = HeavyLatencyMs {
485 pre_phase: pre_phase_ms,
486 parse_changed: t_parse_changed.duration_since(t0).as_secs_f64() * 1000.0,
487 universe_walk: t_universe.duration_since(t_parse_changed).as_secs_f64() * 1000.0,
488 discovery: t_discovery.duration_since(t_universe).as_secs_f64() * 1000.0,
489 parse_discovered: t_parse_discovered.duration_since(t_discovery).as_secs_f64() * 1000.0,
490 tokenization: t_tokenization
491 .duration_since(t_parse_discovered)
492 .as_secs_f64()
493 * 1000.0,
494 graph_build: graph_build_ms,
495 scoring: (t_done.duration_since(t_tokenization).as_secs_f64() * 1000.0 - graph_build_ms)
496 .max(0.0),
497 };
498
499 tracing::debug!(
500 "diffctx heavy: pre_phase {:.3}s, parse_changed {:.3}s, universe {:.3}s, discovery {:.3}s, parse_discovered {:.3}s, tokenization {:.3}s, graph_build {:.3}s, scoring {:.3}s",
501 heavy_latency_ms.pre_phase / 1000.0,
502 heavy_latency_ms.parse_changed / 1000.0,
503 heavy_latency_ms.universe_walk / 1000.0,
504 heavy_latency_ms.discovery / 1000.0,
505 heavy_latency_ms.parse_discovered / 1000.0,
506 heavy_latency_ms.tokenization / 1000.0,
507 heavy_latency_ms.graph_build / 1000.0,
508 heavy_latency_ms.scoring / 1000.0,
509 );
510
511 Ok(ScoredState {
512 root_dir,
513 config,
514 all_fragments,
515 core_ids,
516 core_excerpts,
517 scoring_result,
518 needs,
519 discovery_source,
520 changed_files,
521 deleted_files: deleted_display,
522 renamed_files: renamed_display,
523 lockfile_changes: lockfile_display,
524 ignored_changes: ignored_display,
525 policy_excluded_count: policy_excluded,
526 preferred_revs,
527 commit_message,
528 heavy_latency_ms,
529 })
530}
531
532pub struct SelectionOutcome {
533 pub selected: Vec<Fragment>,
534 pub effective_budget: u32,
535 pub selection_iters: usize,
536 pub stopping_certificate: f64,
537 pub select_ms: f64,
538}
539
540pub fn run_selection(
544 state: &ScoredState,
545 budget_tokens: Option<u32>,
546 tau: f64,
547) -> SelectionOutcome {
548 let t_start = Instant::now();
549 let effective_budget = budget_tokens.unwrap_or_else(|| {
550 let core_tokens: u32 = state
551 .all_fragments
552 .iter()
553 .filter(|f| state.core_ids.contains(&f.id))
554 .map(|f| f.token_count.min(BUDGET.core_token_cap_per_fragment))
555 .sum();
556 let auto = (core_tokens as f64 * BUDGET.auto_multiplier) as u32;
557 auto.clamp(BUDGET.auto_min, BUDGET.auto_max)
558 });
559
560 let selection_result = match state.config.objective {
561 crate::mode::ObjectiveMode::BoltzmannModular => {
562 let beta = crate::utility::calibrate_beta(
563 &state.scoring_result.filtered_fragments,
564 &state.core_ids,
565 &state.scoring_result.rel_scores,
566 effective_budget,
567 crate::config::selection::boltzmann().calibration_tolerance,
568 );
569 tracing::debug!("diffctx: boltzmann beta calibrated to {:.6e}", beta);
570 crate::utility::boltzmann_select(
571 &state.scoring_result.filtered_fragments,
572 &state.core_ids,
573 &state.scoring_result.rel_scores,
574 effective_budget,
575 beta,
576 )
577 }
578 crate::mode::ObjectiveMode::Submodular => {
579 let file_importance =
580 crate::utility::compute_file_importance(&state.scoring_result.filtered_fragments);
581 crate::select::lazy_greedy_select(
582 state.scoring_result.filtered_fragments.clone(),
583 &state.core_ids,
584 &state.scoring_result.rel_scores,
585 &state.needs,
586 effective_budget,
587 tau,
588 Some(&file_importance),
589 Some(&state.core_excerpts),
590 )
591 }
592 };
593
594 let selection_iters = selection_result.greedy_iters;
595 let stopping_certificate = selection_result.stopping_certificate;
596 let mut selected = selection_result.selected;
597
598 postpass::coherence_post_pass(
599 &mut selected,
600 &state.scoring_result.filtered_fragments,
601 &state.scoring_result.graph,
602 effective_budget,
603 );
604
605 postpass::rescue_nontrivial_context(
606 &mut selected,
607 &state.all_fragments,
608 &state.scoring_result.rel_scores,
609 &state.core_ids,
610 effective_budget,
611 );
612
613 let used: u32 = selected.iter().map(|f| f.token_count).sum();
614 let remaining = effective_budget.saturating_sub(used);
615 let mut batch_reader = match CatFileBatch::new(&state.root_dir) {
616 Ok(r) => Some(r),
617 Err(_) => None,
618 };
619 postpass::ensure_changed_files_represented(
620 &mut selected,
621 &state.all_fragments,
622 &state.changed_files,
623 remaining,
624 &state.root_dir,
625 &state.preferred_revs,
626 batch_reader.as_mut(),
627 &state.core_ids,
628 &state.core_excerpts,
629 );
630 if let Some(mut r) = batch_reader {
631 r.close();
632 }
633
634 crate::provenance::maybe_dump(state, &selected);
635
636 let select_ms = t_start.elapsed().as_secs_f64() * 1000.0;
637 SelectionOutcome {
638 selected,
639 effective_budget,
640 selection_iters,
641 stopping_certificate,
642 select_ms,
643 }
644}
645
646pub fn select_with_params(
654 state: &ScoredState,
655 budget_tokens: Option<u32>,
656 tau: f64,
657 no_content: bool,
658) -> DiffContextOutput {
659 let outcome = run_selection(state, budget_tokens, tau);
660 let selected = outcome.selected;
661 let selection_iters = outcome.selection_iters;
662 let stopping_certificate = outcome.stopping_certificate;
663 let select_ms = outcome.select_ms;
664
665 let total_ms = state.heavy_latency_ms.pre_phase
666 + state.heavy_latency_ms.parse_changed
667 + state.heavy_latency_ms.universe_walk
668 + state.heavy_latency_ms.discovery
669 + state.heavy_latency_ms.parse_discovered
670 + state.heavy_latency_ms.tokenization
671 + state.heavy_latency_ms.graph_build
672 + state.heavy_latency_ms.scoring
673 + select_ms;
674
675 let cap_stats = state.scoring_result.graph.cap_stats.clone();
676 let change = render::ChangeSummary {
677 commit_message: state.commit_message.clone(),
678 changed_files: state
679 .changed_files
680 .iter()
681 .map(|p| {
682 p.strip_prefix(&state.root_dir)
683 .unwrap_or(p)
684 .to_string_lossy()
685 .replace('\\', "/")
686 })
687 .collect(),
688 deleted_files: state.deleted_files.clone(),
689 renamed_files: state.renamed_files.clone(),
690 lockfile_changes: state.lockfile_changes.clone(),
691 ignored_changes: state.ignored_changes.clone(),
692 policy_excluded_count: state.policy_excluded_count,
693 };
694 let mut render_core_ids = state.core_ids.clone();
698 render_core_ids.extend(
699 selected
700 .iter()
701 .filter(|f| f.kind == crate::types::FragmentKind::Excerpt)
702 .map(|f| f.id.clone()),
703 );
704
705 let mut output = render::build_diff_context_output(
706 &state.root_dir,
707 &selected,
708 no_content,
709 &render_core_ids,
710 &state.scoring_result.rel_scores,
711 change,
712 );
713 tracing::debug!(
714 "diffctx selection: selection {:.3}s (incl. post-passes), total {:.3}s",
715 select_ms / 1000.0,
716 total_ms / 1000.0,
717 );
718 output.latency = Some(render::LatencyBreakdown {
719 pre_phase_ms: state.heavy_latency_ms.pre_phase,
720 parse_changed_ms: state.heavy_latency_ms.parse_changed,
721 universe_walk_ms: state.heavy_latency_ms.universe_walk,
722 discovery_ms: state.heavy_latency_ms.discovery,
723 parse_discovered_ms: state.heavy_latency_ms.parse_discovered,
724 tokenization_ms: state.heavy_latency_ms.tokenization,
725 graph_build_ms: state.heavy_latency_ms.graph_build,
726 scoring_selection_ms: state.heavy_latency_ms.graph_build
727 + state.heavy_latency_ms.scoring
728 + select_ms,
729 total_ms,
730 scoring_ms: state.heavy_latency_ms.scoring,
731 selection_ms: select_ms,
732 candidate_count: state.scoring_result.filtered_fragments.len(),
733 edge_count: state.scoring_result.graph.edge_count(),
734 greedy_iters: selection_iters,
735 edges_before_cap: cap_stats.edges_before_cap,
736 edges_dropped_by_cap: cap_stats.edges_dropped_by_cap,
737 nodes_capped: cap_stats.nodes_capped,
738 max_out_edges_per_node: cap_stats.max_out_edges_per_node,
739 ppr_truncated: state.scoring_result.ppr_truncated,
740 stopping_certificate,
741 ppr_forward_pushes: state.scoring_result.ppr_forward_pushes,
742 ppr_backward_pushes: state.scoring_result.ppr_backward_pushes,
743 peak_rss_bytes: crate::peak_rss::peak_rss_bytes(),
744 edge_emissions_by_category: cap_stats
745 .emissions_by_category
746 .iter()
747 .map(|&(category, raw, deduped)| (category.as_str(), raw, deduped))
748 .collect(),
749 });
750 output
751}
752
753pub(crate) fn is_secret_path(path: &Path) -> bool {
762 let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
763 return false;
764 };
765 if matches!(
769 name,
770 "id_rsa"
771 | "id_dsa"
772 | "id_ecdsa"
773 | "id_ed25519"
774 | "id_ed25519_sk"
775 | "id_ecdsa_sk"
776 | ".netrc"
777 | "_netrc"
778 | "credentials"
779 | ".npmrc"
780 | ".pypirc"
781 ) {
782 return true;
783 }
784 matches!(
785 path.extension().and_then(|e| e.to_str()),
786 Some("pem" | "key" | "pfx" | "p12" | "keystore" | "jks" | "ppk" | "p8" | "asc")
790 )
791}
792
793fn is_lockfile_path(path: &Path) -> bool {
799 let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
800 return false;
801 };
802 matches!(
803 name,
804 "Cargo.lock"
805 | "package-lock.json"
806 | "npm-shrinkwrap.json"
807 | "yarn.lock"
808 | "pnpm-lock.yaml"
809 | "bun.lock"
810 | "bun.lockb"
811 | "deno.lock"
812 | "Pipfile.lock"
813 | "poetry.lock"
814 | "uv.lock"
815 | "pdm.lock"
816 | "composer.lock"
817 | "Gemfile.lock"
818 | "flake.lock"
819 | "go.sum"
820 | "mix.lock"
821 | "packages.lock.json"
822 | "gradle.lockfile"
823 | "Package.resolved"
824 | "cabal.project.freeze"
825 )
826}
827
828pub(crate) fn rel_path_string(root_dir: &Path, path: &Path) -> Option<String> {
829 crate::paths::display_rel(root_dir, path)
830}
831
832fn resolve_ignored_paths(
838 root_dir: &Path,
839 hunks: &[crate::types::DiffHunk],
840) -> rustc_hash::FxHashMap<String, git::IgnoreSource> {
841 let rel_paths: Vec<String> = hunks
842 .iter()
843 .filter_map(|h| rel_path_string(root_dir, Path::new(&*h.path)))
844 .collect();
845 git::find_ignored_paths_with_source(root_dir, &rel_paths)
846}
847
848pub(crate) fn is_ignored_path(
849 root_dir: &Path,
850 path: &Path,
851 ignored_rel_paths: &rustc_hash::FxHashMap<String, git::IgnoreSource>,
852) -> bool {
853 rel_path_string(root_dir, path)
854 .map(|rel| ignored_rel_paths.contains_key(&rel))
855 .unwrap_or(false)
856}
857
858pub fn raw_diff_text(root_dir: &Path, diff_range: Option<&str>, timeout: u64) -> Result<String> {
863 git::set_git_timeout(timeout);
864 crate::deadline::set_compute_deadline(timeout);
865 let root_dir = resolve_repo_root(root_dir)?;
866 let resolved = git::resolve_duration_range(&root_dir, diff_range)?;
867 let diff_text = git::get_diff_text(&root_dir, resolved.range.as_deref())?;
868 Ok(keep_disclosable_sections(&root_dir, &diff_text))
869}
870
871fn keep_disclosable_sections(root_dir: &Path, diff_text: &str) -> String {
872 let raw: Vec<&str> = diff_text.split_inclusive('\n').collect();
878 let lines: Vec<&str> = raw
879 .iter()
880 .map(|line| line.trim_end_matches('\n').trim_end_matches('\r'))
881 .collect();
882 let sections = split_diff_sections(root_dir, &lines);
883 let rel_paths: Vec<String> = sections
884 .iter()
885 .filter_map(|(path, _)| path.as_deref())
886 .filter_map(|path| rel_path_string(root_dir, path))
887 .collect();
888 let ignored_rel_paths = git::find_ignored_paths_with_source(root_dir, &rel_paths);
889
890 let mut kept: Vec<&str> = Vec::new();
891 for (path, range) in sections {
892 let Some(path) = path else {
896 continue;
897 };
898 if is_secret_path(&path)
899 || is_lockfile_path(&path)
900 || is_ignored_path(root_dir, &path, &ignored_rel_paths)
901 {
902 continue;
903 }
904 kept.extend_from_slice(&raw[range]);
905 }
906 if kept.is_empty() {
907 return String::new();
908 }
909 let mut text = kept.concat();
910 if !text.ends_with('\n') {
911 text.push('\n');
912 }
913 text
914}
915
916type DiffSection = (Option<PathBuf>, std::ops::Range<usize>);
917
918fn split_diff_sections(root_dir: &Path, lines: &[&str]) -> Vec<DiffSection> {
919 let starts: Vec<usize> = lines
920 .iter()
921 .enumerate()
922 .filter(|(_, line)| line.starts_with("diff --git "))
923 .map(|(index, _)| index)
924 .collect();
925 starts
926 .iter()
927 .enumerate()
928 .map(|(nth, &start)| {
929 let end = starts.get(nth + 1).copied().unwrap_or(lines.len());
930 (section_path(root_dir, &lines[start..end]), start..end)
931 })
932 .collect()
933}
934
935fn section_path(root_dir: &Path, section: &[&str]) -> Option<PathBuf> {
936 let mut old_path: Option<PathBuf> = None;
937 let mut new_path: Option<PathBuf> = None;
938 for line in section.iter().take_while(|line| !line.starts_with("@@")) {
939 match git::parse_path_line(line, root_dir) {
940 ("new", path) => new_path = path,
941 ("old", path) => old_path = path,
942 _ => {}
943 }
944 }
945 new_path
946 .or(old_path)
947 .or_else(|| pathless_section(root_dir, section))
948}
949
950fn pathless_section(root_dir: &Path, section: &[&str]) -> Option<PathBuf> {
956 if let Some(quoted) = section
957 .iter()
958 .find_map(|line| line.strip_prefix("rename to "))
959 {
960 return git::resolve_in_repo(root_dir, &git::unquote_c_style(quoted.trim()));
961 }
962 let rest = section.first()?.strip_prefix("diff --git ")?;
963 let rel_path = rest.get("a/".len()..rest.find(" b/")?)?;
964 if rest != format!("a/{rel_path} b/{rel_path}") {
965 return None;
966 }
967 git::resolve_in_repo(root_dir, rel_path)
968}
969
970fn resolve_repo_root(root_dir: &Path) -> Result<PathBuf> {
971 let root_dir = root_dir.canonicalize().unwrap_or_else(|e| {
972 tracing::debug!("canonicalize failed for '{}': {}", root_dir.display(), e);
973 root_dir.to_path_buf()
974 });
975 if !git::is_git_repo(&root_dir) {
976 anyhow::bail!("'{}' is not a git repository", root_dir.display());
977 }
978 Ok(git::find_toplevel(&root_dir).unwrap_or(root_dir))
979}
980
981fn build_diff_context_full(
982 root_dir: &Path,
983 diff_range: Option<&str>,
984 no_content: bool,
985 timeout: u64,
986) -> Result<DiffContextOutput> {
987 git::set_git_timeout(timeout);
988 crate::deadline::set_compute_deadline(timeout);
989 let root_dir = resolve_repo_root(root_dir)?;
990 let resolved = git::resolve_duration_range(&root_dir, diff_range)?;
991 let diff_range = resolved.range.as_deref();
992 let mut hunks = git::parse_diff(&root_dir, diff_range)?;
993 hunks.retain(|h| !is_secret_path(Path::new(&*h.path)));
994 if hunks.is_empty() {
995 let (deleted, renamed) = deletion_rename_displays(&root_dir, diff_range);
996 let mut output = empty_output(&root_dir);
997 output.deleted_files = deleted;
998 output.renamed_files = renamed;
999 return Ok(output);
1000 }
1001 let ignored_rel_paths = resolve_ignored_paths(&root_dir, &hunks);
1002 hunks.retain(|h| !is_ignored_path(&root_dir, Path::new(&*h.path), &ignored_rel_paths));
1003 if hunks.is_empty() {
1004 let (deleted, renamed) = deletion_rename_displays(&root_dir, diff_range);
1005 let mut output = empty_output(&root_dir);
1006 output.deleted_files = deleted;
1007 output.renamed_files = renamed;
1008 return Ok(output);
1009 }
1010 let mut changed_files = git::get_changed_files(&root_dir, diff_range)?;
1011 changed_files
1012 .retain(|f| !is_secret_path(f) && !is_ignored_path(&root_dir, f, &ignored_rel_paths));
1013 if changed_files.is_empty() {
1014 let (deleted, renamed) = deletion_rename_displays(&root_dir, diff_range);
1015 let mut output = empty_output(&root_dir);
1016 output.deleted_files = deleted;
1017 output.renamed_files = renamed;
1018 return Ok(output);
1019 }
1020 let (base_rev, head_rev) = diff_range
1021 .map(git::split_diff_range)
1022 .unwrap_or((None, None));
1023 let preferred_revs = build_preferred_revs(base_rev.as_deref(), head_rev.as_deref());
1024 let mut seen_frag_ids: FxHashSet<FragmentId> = FxHashSet::default();
1025 let mut batch_reader = CatFileBatch::new(&root_dir)?;
1026 let mut all_fragments = process_files_for_fragments(
1027 &changed_files,
1028 &root_dir,
1029 &preferred_revs,
1030 &mut seen_frag_ids,
1031 Some(&mut batch_reader),
1032 true,
1033 );
1034 assign_token_counts(&mut all_fragments);
1035 let mut sig_frags = generate_signature_variants(&all_fragments);
1036 assign_token_counts(&mut sig_frags);
1037 all_fragments.extend(sig_frags);
1038 changed_files.sort();
1039 let core_ids = identify_core_fragments(&hunks, &all_fragments);
1040 let selected = select_full_mode(&all_fragments, &changed_files);
1041 batch_reader.close();
1042 let commit_message = head_rev
1043 .as_deref()
1044 .and_then(|h| git::get_commit_message(&root_dir, h).ok())
1045 .and_then(|m| {
1046 m.lines()
1047 .map(str::trim)
1048 .find(|l| !l.is_empty())
1049 .map(str::to_string)
1050 });
1051 let mut deleted_display: Vec<String> = git::get_deleted_files(&root_dir, diff_range)
1052 .map(|set| {
1053 set.iter()
1054 .map(|p| {
1055 p.strip_prefix(&root_dir)
1056 .unwrap_or(p)
1057 .to_string_lossy()
1058 .replace('\\', "/")
1059 })
1060 .collect()
1061 })
1062 .unwrap_or_default();
1063 deleted_display.sort();
1064 let change = render::ChangeSummary {
1065 commit_message,
1066 changed_files: changed_files
1067 .iter()
1068 .map(|p| {
1069 p.strip_prefix(&root_dir)
1070 .unwrap_or(p)
1071 .to_string_lossy()
1072 .replace('\\', "/")
1073 })
1074 .collect(),
1075 deleted_files: deleted_display,
1076 renamed_files: git::get_rename_pairs(&root_dir, diff_range).unwrap_or_default(),
1077 lockfile_changes: Vec::new(),
1080 ignored_changes: Vec::new(),
1081 policy_excluded_count: 0,
1082 };
1083 Ok(render::build_diff_context_output(
1084 &root_dir,
1085 &selected,
1086 no_content,
1087 &core_ids,
1088 &FxHashMap::default(),
1089 change,
1090 ))
1091}
1092
1093fn deletion_rename_displays(
1094 root_dir: &Path,
1095 diff_range: Option<&str>,
1096) -> (Vec<String>, Vec<(String, String)>) {
1097 let mut deleted: Vec<String> = git::get_deleted_files(root_dir, diff_range)
1098 .map(|set| {
1099 set.iter()
1100 .map(|p| {
1101 p.strip_prefix(root_dir)
1102 .unwrap_or(p)
1103 .to_string_lossy()
1104 .replace('\\', "/")
1105 })
1106 .collect()
1107 })
1108 .unwrap_or_default();
1109 deleted.sort();
1110 let renamed = git::get_rename_pairs(root_dir, diff_range).unwrap_or_default();
1111 (deleted, renamed)
1112}
1113
1114fn empty_scored_state_with_changes(root_dir: PathBuf, diff_range: Option<&str>) -> ScoredState {
1115 let (deleted, renamed) = deletion_rename_displays(&root_dir, diff_range);
1116 let mut state = empty_scored_state(root_dir);
1117 state.deleted_files = deleted;
1118 state.renamed_files = renamed;
1119 state
1120}
1121
1122fn empty_scored_state(root_dir: PathBuf) -> ScoredState {
1123 let config = PipelineConfig::from_mode(ScoringMode::Ego);
1124 ScoredState {
1125 root_dir,
1126 config,
1127 all_fragments: Vec::new(),
1128 core_ids: FxHashSet::default(),
1129 core_excerpts: FxHashMap::default(),
1130 discovery_source: FxHashMap::default(),
1131 lockfile_changes: Vec::new(),
1132 ignored_changes: Vec::new(),
1133 policy_excluded_count: 0,
1134 scoring_result: ScoringResult {
1135 rel_scores: FxHashMap::default(),
1136 filtered_fragments: Vec::new(),
1137 graph: crate::graph::Graph::new(),
1138 graph_build_ms: 0.0,
1139 ppr_truncated: false,
1140 ppr_forward_pushes: 0,
1141 ppr_backward_pushes: 0,
1142 },
1143 needs: Vec::new(),
1144 changed_files: Vec::new(),
1145 deleted_files: Vec::new(),
1146 renamed_files: Vec::new(),
1147 preferred_revs: Vec::new(),
1148 commit_message: None,
1149 heavy_latency_ms: HeavyLatencyMs::default(),
1150 }
1151}
1152
1153fn empty_output(root_dir: &Path) -> DiffContextOutput {
1154 let resolved = root_dir
1155 .canonicalize()
1156 .unwrap_or_else(|_| root_dir.to_path_buf());
1157 let name = resolved
1158 .file_name()
1159 .map(|n| n.to_string_lossy().to_string())
1160 .unwrap_or_else(|| resolved.to_string_lossy().to_string());
1161 DiffContextOutput {
1162 name,
1163 output_type: "diff_context".to_string(),
1164 commit_message: None,
1165 changed_files: Vec::new(),
1166 deleted_files: Vec::new(),
1167 renamed_files: Vec::new(),
1168 lockfile_changes: Vec::new(),
1169 ignored_changes: Vec::new(),
1170 policy_excluded_count: 0,
1171 fragment_count: 0,
1172 fragments: Vec::new(),
1173 latency: None,
1174 }
1175}
1176
1177pub(crate) fn empty_output_from_state(state: &ScoredState) -> DiffContextOutput {
1182 let mut output = empty_output(&state.root_dir);
1183 output.commit_message = state.commit_message.clone();
1184 output.deleted_files = state.deleted_files.clone();
1185 output.renamed_files = state.renamed_files.clone();
1186 output.lockfile_changes = state.lockfile_changes.clone();
1187 output.ignored_changes = state.ignored_changes.clone();
1188 output.policy_excluded_count = state.policy_excluded_count;
1189 output
1190}
1191
1192fn build_preferred_revs(base_rev: Option<&str>, head_rev: Option<&str>) -> Vec<String> {
1193 let mut revs = Vec::new();
1194 if let Some(h) = head_rev {
1195 revs.push(h.to_string());
1196 }
1197 if let Some(b) = base_rev {
1198 if Some(b) != head_rev {
1199 revs.push(b.to_string());
1200 }
1201 }
1202 revs
1203}
1204
1205fn create_discovery(config: &PipelineConfig) -> Box<dyn DiscoveryStrategy> {
1206 match config.discovery {
1207 DiscoveryKind::Ensemble => Box::new(EnsembleDiscovery::new(vec![
1208 Box::new(DefaultDiscovery),
1209 Box::new(TestFileDiscovery),
1210 Box::new(BM25Discovery::new(config.bm25_top_k)),
1211 ])),
1212 DiscoveryKind::Default => Box::new(DefaultDiscovery),
1213 }
1214}
1215
1216fn build_file_cache(candidate_files: &[PathBuf]) -> FxHashMap<PathBuf, String> {
1217 let mut sorted = candidate_files.to_vec();
1222 sorted.sort();
1223 let mut cache: FxHashMap<PathBuf, String> = FxHashMap::default();
1224 let mut cache_bytes = 0usize;
1225 for path in sorted {
1226 if cache_bytes > GRAPH_FILTERING.max_cache_bytes {
1227 break;
1228 }
1229 let Ok(meta) = path.metadata() else { continue };
1230 if meta.len() as usize > LIMITS.max_file_size {
1231 continue;
1232 }
1233 if let Ok(content) = std::fs::read_to_string(&path) {
1234 cache_bytes += content.len();
1235 cache.insert(path, content);
1236 }
1237 }
1238 cache
1239}
1240
1241fn assign_token_counts(fragments: &mut [Fragment]) {
1242 fragments.par_iter_mut().for_each(|frag| {
1243 if frag.token_count == 0 {
1244 frag.token_count = count_tokens(&frag.content) + LIMITS.overhead_per_fragment;
1245 }
1246 });
1247}
1248
1249fn assign_excerpt_token_counts(excerpts: &mut FxHashMap<FragmentId, Fragment>) {
1250 for frag in excerpts.values_mut() {
1251 if frag.token_count == 0 {
1252 frag.token_count = count_tokens(&frag.content) + LIMITS.overhead_per_fragment;
1253 }
1254 }
1255}
1256
1257fn select_full_mode(all_fragments: &[Fragment], changed_files: &[PathBuf]) -> Vec<Fragment> {
1258 let changed_paths: FxHashSet<String> = changed_files
1259 .iter()
1260 .map(|p| p.to_string_lossy().to_string())
1261 .collect();
1262 let mut selected: Vec<Fragment> = all_fragments
1263 .iter()
1264 .filter(|f| changed_paths.contains(f.path()))
1265 .cloned()
1266 .collect();
1267 selected.sort_by(|a, b| {
1268 a.path()
1269 .cmp(b.path())
1270 .then(a.start_line().cmp(&b.start_line()))
1271 });
1272 selected
1273}
1274
1275#[cfg(test)]
1276mod tests {
1277 use super::*;
1278
1279 #[test]
1286 fn a_raw_diff_section_escaping_the_root_resolves_to_nothing() {
1287 let tmp = tempfile::TempDir::new().expect("tempdir");
1288 let base = tmp.path().canonicalize().expect("canonical tempdir");
1289 let root = base.join("repo");
1290 std::fs::create_dir_all(&root).expect("mkdir repo");
1291 std::fs::write(base.join("outside.py"), "secret = 1\n").expect("write outside");
1292 std::fs::write(root.join("inside.py"), "x = 1\n").expect("write inside");
1293
1294 for rel in ["../outside.py", "../missing.py", "sub/../../outside.py"] {
1295 let header = format!("diff --git a/{rel} b/{rel}");
1296 assert!(
1297 section_path(&root, &[header.as_str()]).is_none(),
1298 "escaping section accepted: {rel}"
1299 );
1300 }
1301
1302 let header = "diff --git a/inside.py b/inside.py";
1303 assert!(
1304 section_path(&root, &[header]).is_some_and(|p| p.ends_with("inside.py")),
1305 "an in-repo section was dropped"
1306 );
1307 }
1308}
1309
1310#[cfg(test)]
1311mod secret_path_tests {
1312 use super::is_secret_path;
1313 use std::path::Path;
1314
1315 fn secret(p: &str) -> bool {
1316 is_secret_path(Path::new(p))
1317 }
1318
1319 #[test]
1325 fn private_key_and_credential_shapes_are_excluded() {
1326 for path in [
1327 "home/.ssh/id_rsa",
1328 "home/.ssh/id_ed25519",
1329 "home/.ssh/id_ed25519_sk",
1330 "home/.ssh/id_ecdsa_sk",
1331 "certs/server.pem",
1332 "certs/server.key",
1333 "certs/bundle.pfx",
1334 "certs/bundle.p12",
1335 "android/release.keystore",
1336 "android/release.jks",
1337 "windows/deploy.ppk",
1338 "apple/AuthKey_ABC123.p8",
1339 "gpg/private.asc",
1340 "home/.netrc",
1341 "home/_netrc",
1342 "aws/credentials",
1343 "home/.npmrc",
1344 "home/.pypirc",
1345 ] {
1346 assert!(secret(path), "not excluded: {path}");
1347 }
1348 }
1349
1350 #[test]
1353 fn public_material_is_not_excluded() {
1354 for path in [
1355 "home/.ssh/id_rsa.pub",
1356 "home/.ssh/id_ed25519.pub",
1357 "home/.ssh/authorized_keys",
1358 "certs/server.crt",
1359 ] {
1360 assert!(!secret(path), "wrongly excluded: {path}");
1361 }
1362 }
1363
1364 #[test]
1368 fn env_files_remain_visible_by_design() {
1369 assert!(!secret(".env"));
1370 assert!(!secret("config/.env.production"));
1371 }
1372
1373 #[test]
1376 fn ordinary_files_are_untouched() {
1377 for path in [
1378 "src/keyboard.rs",
1379 "src/credentials_form.tsx",
1380 "docs/pemphigus.md",
1381 ] {
1382 assert!(!secret(path), "wrongly excluded: {path}");
1383 }
1384 }
1385}