Skip to main content

_diffctx/edges/structural/
sibling.rs

1use std::path::Path;
2
3use rustc_hash::{FxHashMap, FxHashSet};
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        // A path belongs to exactly one directory, so one set of already-seen
18        // paths is enough to dedupe every bucket. This used to be
19        // `Vec::contains` against the bucket, i.e. a scan of the directory's
20        // whole file list per fragment — quadratic in exactly the shape this
21        // builder exists for (thousands of files in one directory), while
22        // producing the same buckets in the same first-seen order.
23        let mut seen: FxHashSet<&str> = FxHashSet::default();
24        for f in fragments {
25            let path_str = f.path();
26            if !seen.insert(path_str) {
27                continue;
28            }
29            let dir = Path::new(path_str)
30                .parent()
31                .map(|p| p.to_string_lossy().to_string())
32                .unwrap_or_default();
33            by_dir.entry(dir).or_default().push(path_str);
34        }
35        by_dir
36    }
37
38    fn build_file_representative_map(
39        &self,
40        fragments: &[Fragment],
41    ) -> FxHashMap<String, FragmentId> {
42        super::super::base::file_representatives(fragments)
43    }
44}
45
46impl EdgeBuilder for SiblingEdgeBuilder {
47    fn build(&self, fragments: &[Fragment], _repo_root: Option<&Path>) -> EdgeDict {
48        let weight = EDGE_WEIGHTS["sibling"].forward;
49        let reverse_factor = EDGE_WEIGHTS["sibling"].reverse_factor;
50
51        let by_dir = self.group_files_by_dir(fragments);
52        let file_to_rep = self.build_file_representative_map(fragments);
53
54        let mut edges: EdgeDict = FxHashMap::default();
55
56        for (_dir, files) in &by_dir {
57            let mut file_list: Vec<&str> = files.clone();
58            file_list.sort_unstable();
59            if file_list.len() > SIBLING.max_files_per_dir {
60                file_list.truncate(SIBLING.max_files_per_dir);
61            }
62            if file_list.len() < 2 {
63                continue;
64            }
65
66            for i in 0..file_list.len() {
67                for j in (i + 1)..file_list.len() {
68                    if let (Some(f1_id), Some(f2_id)) =
69                        (file_to_rep.get(file_list[i]), file_to_rep.get(file_list[j]))
70                    {
71                        add_edge(&mut edges, f1_id, f2_id, weight, reverse_factor);
72                    }
73                }
74            }
75        }
76
77        edges
78    }
79
80    fn category_label(&self) -> Option<&str> {
81        Some("sibling")
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use crate::types::FragmentKind;
89    use rustc_hash::FxHashSet as Set;
90    use std::sync::Arc;
91
92    fn frag(path: &str, start: u32) -> Fragment {
93        Fragment {
94            id: crate::types::FragmentId::new(Arc::from(path), start, start + 5),
95            kind: FragmentKind::Function,
96            content: Arc::from(""),
97            identifiers: Set::default(),
98            token_count: 10,
99            symbol_name: None,
100        }
101    }
102
103    /// Grouping deduped files with `Vec::contains` against the bucket, so every
104    /// fragment scanned its directory's whole file list. The buckets it produces
105    /// are what matters: one entry per file, in first-seen order, regardless of
106    /// how many fragments each file contributes.
107    #[test]
108    fn each_file_appears_once_per_directory_in_first_seen_order() {
109        // Several fragments per file, files interleaved across two directories.
110        let fragments = vec![
111            frag("src/b.rs", 1),
112            frag("src/a.rs", 1),
113            frag("src/b.rs", 20),
114            frag("lib/c.rs", 1),
115            frag("src/a.rs", 40),
116            frag("lib/c.rs", 30),
117        ];
118
119        let by_dir = SiblingEdgeBuilder.group_files_by_dir(&fragments);
120
121        assert_eq!(
122            by_dir.get("src").map(Vec::as_slice),
123            Some(&["src/b.rs", "src/a.rs"][..])
124        );
125        assert_eq!(
126            by_dir.get("lib").map(Vec::as_slice),
127            Some(&["lib/c.rs"][..])
128        );
129    }
130
131    /// The per-directory pair loop is quadratic by nature, so the cap is what
132    /// keeps a flat thousand-file directory from emitting a near-dense block.
133    #[test]
134    fn a_directory_over_the_cap_emits_only_the_capped_pairs() {
135        let over = SIBLING.max_files_per_dir + 40;
136        let fragments: Vec<Fragment> = (0..over)
137            .map(|i| frag(&format!("src/f{i:04}.rs"), 1))
138            .collect();
139
140        let edges = SiblingEdgeBuilder.build(&fragments, None);
141        let k = SIBLING.max_files_per_dir;
142        // Every kept pair contributes a forward and a reverse edge.
143        assert_eq!(edges.len(), k * (k - 1));
144    }
145}