_diffctx/edges/structural/
sibling.rs1use std::path::Path;
2
3use rustc_hash::FxHashMap;
4
5use crate::config::limits::SIBLING;
6use crate::config::weights::EDGE_WEIGHTS;
7use crate::types::{Fragment, FragmentId};
8
9use super::super::EdgeDict;
10use super::super::base::{EdgeBuilder, add_edge};
11
12pub struct SiblingEdgeBuilder;
13
14impl SiblingEdgeBuilder {
15 fn group_files_by_dir<'a>(&self, fragments: &'a [Fragment]) -> FxHashMap<String, Vec<&'a str>> {
16 let mut by_dir: FxHashMap<String, Vec<&str>> = FxHashMap::default();
17 for f in fragments {
18 let path = Path::new(f.path());
19 let dir = path
20 .parent()
21 .map(|p| p.to_string_lossy().to_string())
22 .unwrap_or_default();
23 let path_str = f.path();
24 let files = by_dir.entry(dir).or_default();
25 if !files.contains(&path_str) {
26 files.push(path_str);
27 }
28 }
29 by_dir
30 }
31
32 fn build_file_representative_map(
33 &self,
34 fragments: &[Fragment],
35 ) -> FxHashMap<String, FragmentId> {
36 let mut file_to_rep: FxHashMap<String, FragmentId> = FxHashMap::default();
37 let mut file_to_token_count: FxHashMap<String, u32> = FxHashMap::default();
38
39 for f in fragments {
40 let path = f.path().to_string();
41 let existing_count = file_to_token_count.get(&path).copied().unwrap_or(0);
42 if !file_to_rep.contains_key(&path) || f.token_count > existing_count {
43 file_to_rep.insert(path.clone(), f.id.clone());
44 file_to_token_count.insert(path, f.token_count);
45 }
46 }
47
48 file_to_rep
49 }
50}
51
52impl EdgeBuilder for SiblingEdgeBuilder {
53 fn build(&self, fragments: &[Fragment], _repo_root: Option<&Path>) -> EdgeDict {
54 let weight = EDGE_WEIGHTS["sibling"].forward;
55 let reverse_factor = EDGE_WEIGHTS["sibling"].reverse_factor;
56
57 let by_dir = self.group_files_by_dir(fragments);
58 let file_to_rep = self.build_file_representative_map(fragments);
59
60 let mut edges: EdgeDict = FxHashMap::default();
61
62 for (_dir, files) in &by_dir {
63 let mut file_list: Vec<&str> = files.clone();
64 file_list.sort_unstable();
65 if file_list.len() > SIBLING.max_files_per_dir {
66 file_list.truncate(SIBLING.max_files_per_dir);
67 }
68 if file_list.len() < 2 {
69 continue;
70 }
71
72 for i in 0..file_list.len() {
73 for j in (i + 1)..file_list.len() {
74 if let (Some(f1_id), Some(f2_id)) =
75 (file_to_rep.get(file_list[i]), file_to_rep.get(file_list[j]))
76 {
77 add_edge(&mut edges, f1_id, f2_id, weight, reverse_factor);
78 }
79 }
80 }
81 }
82
83 edges
84 }
85
86 fn category_label(&self) -> Option<&str> {
87 Some("sibling")
88 }
89}