1use rustc_hash::{FxHashMap, FxHashSet};
2
3use crate::types::{DiffHunk, Fragment, FragmentId, FragmentKind};
4
5fn kind_priority(kind: FragmentKind) -> u8 {
6 if kind.is_semantic() { 0 } else { 1 }
7}
8
9fn find_core_for_hunk(frags: &[&Fragment], h_start: u32, h_end: u32) -> FxHashSet<FragmentId> {
10 let mut core = FxHashSet::default();
11
12 let covering: Vec<&Fragment> = frags
13 .iter()
14 .copied()
15 .filter(|f| f.start_line() <= h_start && h_end <= f.end_line())
16 .collect();
17 if !covering.is_empty() {
18 let best = covering
19 .iter()
20 .min_by(|a, b| {
21 let ka = kind_priority(a.kind);
22 let kb = kind_priority(b.kind);
23 ka.cmp(&kb).then(a.line_count().cmp(&b.line_count()))
24 })
25 .unwrap();
26 core.insert(best.id.clone());
27 return core;
28 }
29
30 let overlapping: Vec<&Fragment> = frags
31 .iter()
32 .copied()
33 .filter(|f| f.start_line() <= h_end && f.end_line() >= h_start)
34 .collect();
35 if !overlapping.is_empty() {
36 for f in &overlapping {
37 core.insert(f.id.clone());
38 }
39 return core;
40 }
41
42 let before: Vec<&Fragment> = frags
43 .iter()
44 .copied()
45 .filter(|f| f.end_line() < h_start)
46 .collect();
47 let after: Vec<&Fragment> = frags
48 .iter()
49 .copied()
50 .filter(|f| f.start_line() > h_end)
51 .collect();
52 if let Some(nearest_before) = before.iter().max_by_key(|f| f.end_line()) {
53 core.insert(nearest_before.id.clone());
54 }
55 if let Some(nearest_after) = after.iter().min_by_key(|f| f.start_line()) {
56 core.insert(nearest_after.id.clone());
57 }
58
59 core
60}
61
62fn add_container_headers(
63 core_ids: &mut FxHashSet<FragmentId>,
64 frags_by_path: &FxHashMap<&str, Vec<&Fragment>>,
65) {
66 let core_paths: FxHashSet<&str> = core_ids.iter().map(|fid| fid.path.as_ref()).collect();
67 let mut headers_to_add = Vec::new();
68
69 for &path in &core_paths {
70 if let Some(frags) = frags_by_path.get(path) {
71 for frag in frags {
72 if !frag.kind.is_container() || core_ids.contains(&frag.id) {
73 continue;
74 }
75 let contains_core = core_ids.iter().any(|core_id| {
76 core_id.path.as_ref() == path
77 && frag.start_line() <= core_id.start_line
78 && core_id.end_line <= frag.end_line()
79 });
80 if contains_core {
81 headers_to_add.push(frag.id.clone());
82 }
83 }
84 }
85 }
86
87 for h in headers_to_add {
88 core_ids.insert(h);
89 }
90}
91
92pub fn identify_core_fragments(
93 hunks: &[DiffHunk],
94 all_fragments: &[Fragment],
95) -> FxHashSet<FragmentId> {
96 let mut frags_by_path: FxHashMap<&str, Vec<&Fragment>> = FxHashMap::default();
97 for frag in all_fragments {
98 frags_by_path.entry(frag.path()).or_default().push(frag);
99 }
100
101 let mut core_ids = FxHashSet::default();
102 for h in hunks {
103 if let Some(frags) = frags_by_path.get(h.path.as_ref()) {
104 let (h_start, h_end) = h.core_selection_range();
105 core_ids.extend(find_core_for_hunk(frags, h_start, h_end));
106 }
107 }
108
109 add_container_headers(&mut core_ids, &frags_by_path);
110 core_ids
111}
112
113fn map_hunks_to_fragments(
114 hunks: &[DiffHunk],
115 core_ids: &FxHashSet<FragmentId>,
116 all_fragments: &[Fragment],
117) -> FxHashMap<FragmentId, f64> {
118 let mut result: FxHashMap<FragmentId, f64> = FxHashMap::default();
119 for h in hunks {
120 let (h_start, h_end) = h.core_selection_range();
121 let hunk_size = (h_end as i64 - h_start as i64 + 1).max(1) as f64;
122 for frag in all_fragments {
123 if !core_ids.contains(&frag.id) || frag.path() != h.path.as_ref() {
124 continue;
125 }
126 if frag.start_line() <= h_end && frag.end_line() >= h_start {
127 *result.entry(frag.id.clone()).or_insert(0.0) += hunk_size;
128 }
129 }
130 }
131 result
132}
133
134fn add_container_weights(
135 frag_hunk_lines: &mut FxHashMap<FragmentId, f64>,
136 core_ids: &FxHashSet<FragmentId>,
137 all_fragments: &[Fragment],
138) {
139 let mut to_add = Vec::new();
140 for frag in all_fragments {
141 if !core_ids.contains(&frag.id) || frag_hunk_lines.contains_key(&frag.id) {
142 continue;
143 }
144 if !frag.kind.is_container() {
145 continue;
146 }
147 let contained_weight: f64 = frag_hunk_lines
148 .iter()
149 .filter(|(fid, _)| {
150 fid.path.as_ref() == frag.path()
151 && frag.start_line() <= fid.start_line
152 && fid.end_line <= frag.end_line()
153 })
154 .map(|(_, w)| *w)
155 .sum();
156 if contained_weight > 0.0 {
157 to_add.push((frag.id.clone(), contained_weight));
158 }
159 }
160 for (id, w) in to_add {
161 frag_hunk_lines.insert(id, w);
162 }
163}
164
165fn best_hunk_size_for_path(hunks: &[DiffHunk], path: &str) -> u32 {
166 let mut best = 0u32;
167 for h in hunks {
168 if h.path.as_ref() == path {
169 let (h_start, h_end) = h.core_selection_range();
170 let size = h_end.saturating_sub(h_start) + 1;
171 best = best.max(size);
172 }
173 }
174 best
175}
176
177fn fill_missing_core_weights(
178 frag_hunk_lines: &mut FxHashMap<FragmentId, f64>,
179 core_ids: &FxHashSet<FragmentId>,
180 hunks: &[DiffHunk],
181) {
182 let missing: Vec<FragmentId> = core_ids
183 .iter()
184 .filter(|fid| !frag_hunk_lines.contains_key(*fid))
185 .cloned()
186 .collect();
187 for fid in missing {
188 let best = best_hunk_size_for_path(hunks, fid.path.as_ref());
189 if best > 0 {
190 frag_hunk_lines.insert(fid, best as f64);
191 }
192 }
193}
194
195pub fn compute_seed_weights(
196 hunks: &[DiffHunk],
197 core_ids: &FxHashSet<FragmentId>,
198 all_fragments: &[Fragment],
199) -> FxHashMap<FragmentId, f64> {
200 let mut frag_hunk_lines = map_hunks_to_fragments(hunks, core_ids, all_fragments);
201 if frag_hunk_lines.is_empty() {
202 return FxHashMap::default();
203 }
204
205 add_container_weights(&mut frag_hunk_lines, core_ids, all_fragments);
206 fill_missing_core_weights(&mut frag_hunk_lines, core_ids, hunks);
207
208 frag_hunk_lines
209}