1use std::collections::{BTreeMap, BTreeSet, HashSet};
2
3use roaring::RoaringBitmap;
4
5use crate::model::FileDiff;
6use crate::testmap::TestMapIndex;
7
8#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
9pub struct ImpactResult {
10 pub selected_tests: Vec<String>,
11 pub smoke_tests: Vec<String>,
12 pub missing_mappings: Vec<String>,
13 pub stale: bool,
14 pub confidence: f64,
15 pub escalate_full_suite: bool,
16}
17
18#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
19pub struct PlannedTest {
20 pub id: String,
21 pub name: String,
22 pub estimated_overlap_lines: u64,
23 pub marginal_gain_lines: u64,
24}
25
26#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
27pub struct UncoveredBlock {
28 pub file: String,
29 pub start_line: u32,
30 pub end_line: u32,
31}
32
33#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default, PartialEq)]
34pub struct ImpactPlan {
35 pub changed_lines_total: u64,
36 pub changed_lines_covered_by_plan: u64,
37 pub plan_coverage_pct: f64,
38 pub tests: Vec<PlannedTest>,
39 pub uncovered_blocks: Vec<UncoveredBlock>,
40 pub next_command: String,
41}
42
43pub fn select_impacted_tests(index: &TestMapIndex, diffs: &[FileDiff]) -> ImpactResult {
44 let mut tests: BTreeSet<String> = BTreeSet::new();
45 let mut missing = BTreeSet::new();
46
47 for diff in diffs {
48 let mut matched = false;
49 if let Some(candidates) = index.file_to_tests.get(&diff.path) {
50 for t in candidates {
51 tests.insert(t.clone());
52 }
53 matched = true;
54 }
55
56 if let Some(old_path) = diff.old_path.as_deref() {
57 if let Some(candidates) = index.file_to_tests.get(old_path) {
58 for t in candidates {
59 tests.insert(t.clone());
60 }
61 matched = true;
62 }
63 }
64
65 if !matched {
66 missing.insert(diff.path.clone());
67 }
68 }
69
70 ImpactResult {
71 selected_tests: tests.into_iter().collect(),
72 smoke_tests: Vec::new(),
73 missing_mappings: missing.into_iter().collect(),
74 stale: false,
75 confidence: 1.0,
76 escalate_full_suite: false,
77 }
78}
79
80pub fn plan_impacted_tests(
81 index: &TestMapIndex,
82 diffs: &[FileDiff],
83 max_tests: usize,
84 target_coverage: f64,
85) -> ImpactPlan {
86 let mut plan = ImpactPlan {
87 next_command: "covy impact run --plan plan.json -- <your-test-command-template>"
88 .to_string(),
89 ..Default::default()
90 };
91
92 let file_to_idx: BTreeMap<&str, usize> = index
93 .file_index
94 .iter()
95 .enumerate()
96 .map(|(i, f)| (f.as_str(), i))
97 .collect();
98
99 let mut mapped_remaining: BTreeMap<usize, RoaringBitmap> = BTreeMap::new();
100 let mut mapped_file_names: BTreeMap<usize, String> = BTreeMap::new();
101 let mut unmapped_remaining: BTreeMap<String, RoaringBitmap> = BTreeMap::new();
102
103 for diff in diffs {
104 let changed = diff.changed_lines.clone();
105 if changed.is_empty() {
106 continue;
107 }
108
109 let mapped_idx = file_to_idx.get(diff.path.as_str()).copied().or_else(|| {
110 diff.old_path
111 .as_deref()
112 .and_then(|p| file_to_idx.get(p).copied())
113 });
114
115 if let Some(file_idx) = mapped_idx {
116 mapped_file_names
117 .entry(file_idx)
118 .or_insert_with(|| index.file_index[file_idx].clone());
119 mapped_remaining
120 .entry(file_idx)
121 .or_insert_with(RoaringBitmap::new)
122 .extend(changed.iter());
123 } else {
124 unmapped_remaining
125 .entry(diff.path.clone())
126 .or_insert_with(RoaringBitmap::new)
127 .extend(changed.iter());
128 }
129 }
130
131 plan.changed_lines_total =
132 total_bitmap_lines(&mapped_remaining) + total_bitmap_lines_by_name(&unmapped_remaining);
133 if plan.changed_lines_total == 0 {
134 plan.plan_coverage_pct = 1.0;
135 return plan;
136 }
137
138 let original_mapped = mapped_remaining.clone();
139
140 let mut test_rows: Vec<Vec<RoaringBitmap>> = Vec::new();
141 for row in &index.coverage {
142 let mut bitmap_row = Vec::with_capacity(row.len());
143 for lines in row {
144 let mut bm = RoaringBitmap::new();
145 for line in lines {
146 bm.insert(*line);
147 }
148 bitmap_row.push(bm);
149 }
150 test_rows.push(bitmap_row);
151 }
152
153 let max_index_tests = index.tests.len().min(test_rows.len());
154 let mut selected: HashSet<usize> = HashSet::new();
155
156 while plan.tests.len() < max_tests && selected.len() < max_index_tests {
157 let mut best: Option<(usize, u64, u64, String)> = None; for test_idx in 0..max_index_tests {
160 if selected.contains(&test_idx) {
161 continue;
162 }
163 let gain = test_gain_against_remaining(test_rows.get(test_idx), &mapped_remaining);
164 if gain == 0 {
165 continue;
166 }
167 let overlap = test_gain_against_remaining(test_rows.get(test_idx), &original_mapped);
168 let id = index.tests[test_idx].clone();
169
170 best = match best {
171 None => Some((test_idx, gain, overlap, id)),
172 Some((best_idx, best_gain, best_overlap, best_id)) => {
173 if gain > best_gain
174 || (gain == best_gain && overlap > best_overlap)
175 || (gain == best_gain && overlap == best_overlap && id < best_id)
176 {
177 Some((test_idx, gain, overlap, id))
178 } else {
179 Some((best_idx, best_gain, best_overlap, best_id))
180 }
181 }
182 };
183 }
184
185 let Some((winner_idx, winner_gain, winner_overlap, winner_id)) = best else {
186 break;
187 };
188
189 selected.insert(winner_idx);
190 subtract_test_from_remaining(test_rows.get(winner_idx), &mut mapped_remaining);
191
192 plan.tests.push(PlannedTest {
193 id: winner_id.clone(),
194 name: winner_id,
195 estimated_overlap_lines: winner_overlap,
196 marginal_gain_lines: winner_gain,
197 });
198
199 let remaining_total =
200 total_bitmap_lines(&mapped_remaining) + total_bitmap_lines_by_name(&unmapped_remaining);
201 plan.changed_lines_covered_by_plan =
202 plan.changed_lines_total.saturating_sub(remaining_total);
203 plan.plan_coverage_pct =
204 plan.changed_lines_covered_by_plan as f64 / plan.changed_lines_total as f64;
205
206 if plan.plan_coverage_pct >= target_coverage {
207 break;
208 }
209 }
210
211 if plan.changed_lines_total > 0 {
212 let remaining_total =
213 total_bitmap_lines(&mapped_remaining) + total_bitmap_lines_by_name(&unmapped_remaining);
214 plan.changed_lines_covered_by_plan =
215 plan.changed_lines_total.saturating_sub(remaining_total);
216 plan.plan_coverage_pct =
217 plan.changed_lines_covered_by_plan as f64 / plan.changed_lines_total as f64;
218 }
219
220 plan.uncovered_blocks =
221 build_uncovered_blocks(&mapped_remaining, &mapped_file_names, &unmapped_remaining);
222 plan
223}
224
225fn total_bitmap_lines(map: &BTreeMap<usize, RoaringBitmap>) -> u64 {
226 map.values().map(|b| b.len()).sum()
227}
228
229fn total_bitmap_lines_by_name(map: &BTreeMap<String, RoaringBitmap>) -> u64 {
230 map.values().map(|b| b.len()).sum()
231}
232
233fn test_gain_against_remaining(
234 test_row: Option<&Vec<RoaringBitmap>>,
235 mapped_remaining: &BTreeMap<usize, RoaringBitmap>,
236) -> u64 {
237 let Some(test_row) = test_row else {
238 return 0;
239 };
240 let mut gain = 0u64;
241 for (file_idx, remaining) in mapped_remaining {
242 if let Some(test_lines) = test_row.get(*file_idx) {
243 gain += (&remaining.clone() & test_lines).len();
244 }
245 }
246 gain
247}
248
249fn subtract_test_from_remaining(
250 test_row: Option<&Vec<RoaringBitmap>>,
251 mapped_remaining: &mut BTreeMap<usize, RoaringBitmap>,
252) {
253 let Some(test_row) = test_row else {
254 return;
255 };
256
257 let keys: Vec<usize> = mapped_remaining.keys().copied().collect();
258 for file_idx in keys {
259 if let Some(test_lines) = test_row.get(file_idx) {
260 if let Some(remaining) = mapped_remaining.get_mut(&file_idx) {
261 *remaining -= test_lines;
262 if remaining.is_empty() {
263 mapped_remaining.remove(&file_idx);
264 }
265 }
266 }
267 }
268}
269
270fn build_uncovered_blocks(
271 mapped_remaining: &BTreeMap<usize, RoaringBitmap>,
272 mapped_file_names: &BTreeMap<usize, String>,
273 unmapped_remaining: &BTreeMap<String, RoaringBitmap>,
274) -> Vec<UncoveredBlock> {
275 let mut by_file: BTreeMap<String, RoaringBitmap> = BTreeMap::new();
276
277 for (file_idx, lines) in mapped_remaining {
278 if let Some(name) = mapped_file_names.get(file_idx) {
279 by_file.insert(name.clone(), lines.clone());
280 }
281 }
282 for (file, lines) in unmapped_remaining {
283 by_file.insert(file.clone(), lines.clone());
284 }
285
286 let mut blocks = Vec::new();
287 for (file, lines) in by_file {
288 let vec_lines: Vec<u32> = lines.iter().collect();
289 if vec_lines.is_empty() {
290 continue;
291 }
292
293 let mut start = vec_lines[0];
294 let mut end = vec_lines[0];
295 for line in vec_lines.iter().skip(1) {
296 if *line == end + 1 {
297 end = *line;
298 } else {
299 blocks.push(UncoveredBlock {
300 file: file.clone(),
301 start_line: start,
302 end_line: end,
303 });
304 start = *line;
305 end = *line;
306 }
307 }
308 blocks.push(UncoveredBlock {
309 file,
310 start_line: start,
311 end_line: end,
312 });
313 }
314
315 blocks
316}
317
318#[cfg(test)]
319mod tests {
320 use super::*;
321 use crate::model::{DiffStatus, FileDiff};
322
323 fn diff_with_lines(path: &str, lines: &[u32]) -> FileDiff {
324 let mut changed_lines = RoaringBitmap::new();
325 for line in lines {
326 changed_lines.insert(*line);
327 }
328 FileDiff {
329 path: path.to_string(),
330 old_path: None,
331 status: DiffStatus::Modified,
332 changed_lines,
333 }
334 }
335
336 fn basic_v2_map() -> TestMapIndex {
337 let mut map = TestMapIndex::default();
338 map.tests = vec!["t1".to_string(), "t2".to_string(), "t3".to_string()];
339 map.file_index = vec!["src/a.rs".to_string(), "src/b.rs".to_string()];
340 map.coverage = vec![
341 vec![vec![10, 11], vec![]],
342 vec![vec![11], vec![20, 21]],
343 vec![vec![10], vec![20]],
344 ];
345 map
346 }
347
348 #[test]
349 fn test_select_impacted_tests_from_inverse_index() {
350 let mut map = TestMapIndex::default();
351 map.file_to_tests
352 .entry("src/a.rs".to_string())
353 .or_default()
354 .insert("tests::a".to_string());
355
356 let result = select_impacted_tests(&map, &[diff_with_lines("src/a.rs", &[])]);
357 assert_eq!(result.selected_tests, vec!["tests::a".to_string()]);
358 assert!(result.missing_mappings.is_empty());
359 }
360
361 #[test]
362 fn test_select_impacted_tests_missing_mapping() {
363 let map = TestMapIndex::default();
364 let result = select_impacted_tests(&map, &[diff_with_lines("src/missing.rs", &[])]);
365 assert!(result.selected_tests.is_empty());
366 assert_eq!(result.missing_mappings, vec!["src/missing.rs".to_string()]);
367 }
368
369 #[test]
370 fn test_plan_impacted_tests_greedy_and_coverage() {
371 let map = basic_v2_map();
372 let diffs = vec![
373 diff_with_lines("src/a.rs", &[10, 11, 12]),
374 diff_with_lines("src/b.rs", &[20, 21]),
375 ];
376
377 let plan = plan_impacted_tests(&map, &diffs, 2, 0.8);
378 assert_eq!(plan.changed_lines_total, 5);
379 assert_eq!(plan.tests.len(), 2);
380 assert_eq!(plan.tests[0].id, "t2");
381 assert!(plan.plan_coverage_pct >= 0.8);
382 }
383
384 #[test]
385 fn test_plan_impacted_tests_deterministic_tie_break() {
386 let mut map = TestMapIndex::default();
387 map.tests = vec!["aaa".to_string(), "bbb".to_string()];
388 map.file_index = vec!["src/a.rs".to_string()];
389 map.coverage = vec![vec![vec![1]], vec![vec![1]]];
390 let diffs = vec![diff_with_lines("src/a.rs", &[1])];
391
392 let plan = plan_impacted_tests(&map, &diffs, 1, 1.0);
393 assert_eq!(plan.tests[0].id, "aaa");
394 }
395
396 #[test]
397 fn test_uncovered_blocks_compact_ranges() {
398 let mut map = TestMapIndex::default();
399 map.tests = vec!["t1".to_string()];
400 map.file_index = vec!["src/a.rs".to_string()];
401 map.coverage = vec![vec![vec![1, 2]]];
402 let diffs = vec![diff_with_lines("src/a.rs", &[1, 2, 3, 5])];
403
404 let plan = plan_impacted_tests(&map, &diffs, 1, 1.0);
405 assert_eq!(plan.uncovered_blocks.len(), 2);
406 assert_eq!(plan.uncovered_blocks[0].start_line, 3);
407 assert_eq!(plan.uncovered_blocks[0].end_line, 3);
408 assert_eq!(plan.uncovered_blocks[1].start_line, 5);
409 }
410}