Skip to main content

_diffctx/edges/structural/
sibling.rs

1use 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}