lean_ctx/core/
pop_pruning.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use crate::core::task_relevance::RelevanceScore;
4
5#[derive(Debug, Clone)]
6pub struct PopDecision {
7    pub included_modules: Vec<String>,
8    pub excluded_modules: Vec<ExcludedModule>,
9}
10
11#[derive(Debug, Clone)]
12pub struct ExcludedModule {
13    pub module: String,
14    pub candidate_files: usize,
15    pub max_relevance: f64,
16    pub reason: String,
17}
18
19pub fn decide_for_candidates(
20    task: &str,
21    project_root: &str,
22    candidates: &[&RelevanceScore],
23) -> PopDecision {
24    let task_l = task.to_lowercase();
25
26    let mut module_scores: BTreeMap<String, (usize, f64)> = BTreeMap::new(); // (count, max_score)
27    for c in candidates {
28        let m = module_for_path(&c.path, project_root);
29        let e = module_scores.entry(m).or_insert((0, 0.0));
30        e.0 += 1;
31        e.1 = e.1.max(c.score);
32    }
33
34    let mut include: BTreeSet<String> = BTreeSet::new();
35    for m in module_scores.keys() {
36        if module_explicitly_mentioned(&task_l, m) {
37            include.insert(m.clone());
38        }
39    }
40
41    if include.is_empty() {
42        for (m, (_n, max)) in &module_scores {
43            if *max >= 0.7 {
44                include.insert(m.clone());
45            }
46        }
47    }
48
49    let mut excluded = Vec::new();
50    if !include.is_empty() {
51        for (m, (count, max)) in &module_scores {
52            if include.contains(m) {
53                continue;
54            }
55            if *max >= 0.25 {
56                continue;
57            }
58            if *count <= 1 {
59                continue;
60            }
61            excluded.push(ExcludedModule {
62                module: m.clone(),
63                candidate_files: *count,
64                max_relevance: *max,
65                reason: format!("not mentioned by task, max_relevance={:.2} (<0.25)", max),
66            });
67        }
68    }
69
70    PopDecision {
71        included_modules: include.into_iter().collect(),
72        excluded_modules: excluded,
73    }
74}
75
76pub fn filter_candidates_by_pop<'a>(
77    project_root: &str,
78    candidates: &'a [&RelevanceScore],
79    pop: &PopDecision,
80) -> Vec<&'a RelevanceScore> {
81    if pop.excluded_modules.is_empty() {
82        return candidates.to_vec();
83    }
84    let excluded: BTreeSet<&str> = pop
85        .excluded_modules
86        .iter()
87        .map(|e| e.module.as_str())
88        .collect();
89    candidates
90        .iter()
91        .copied()
92        .filter(|c| {
93            let m = module_for_path(&c.path, project_root);
94            !excluded.contains(m.as_str())
95        })
96        .collect()
97}
98
99pub fn module_for_path(path: &str, project_root: &str) -> String {
100    let rel = path
101        .strip_prefix(project_root)
102        .unwrap_or(path)
103        .trim_start_matches('/')
104        .trim_start_matches('\\');
105    rel.split('/')
106        .next()
107        .filter(|s| !s.is_empty())
108        .unwrap_or(".")
109        .to_string()
110}
111
112fn module_explicitly_mentioned(task_l: &str, module: &str) -> bool {
113    if task_l.contains(module) {
114        return true;
115    }
116    match module {
117        "rust" => {
118            task_l.contains("cargo")
119                || task_l.contains("clippy")
120                || task_l.contains("rust")
121                || task_l.contains("ctx_")
122                || task_l.contains("mcp")
123        }
124        "website" => {
125            task_l.contains("website")
126                || task_l.contains("docs")
127                || task_l.contains("astro")
128                || task_l.contains("tailwind")
129                || task_l.contains("gitlab pages")
130        }
131        "packages" => {
132            task_l.contains("vscode")
133                || task_l.contains("chrome")
134                || task_l.contains("extension")
135                || task_l.contains("npm")
136                || task_l.contains("node")
137        }
138        "cloud-infra" => task_l.contains("docker") || task_l.contains("infra"),
139        _ => false,
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn pop_excludes_unmentioned_module() {
149        let root = "/repo";
150        let task = "fix rust bug in ctx_read";
151        let candidates = [
152            RelevanceScore {
153                path: "/repo/rust/src/tools/ctx_read.rs".to_string(),
154                score: 0.9,
155                recommended_mode: "full",
156            },
157            RelevanceScore {
158                path: "/repo/website/src/index.astro".to_string(),
159                score: 0.1,
160                recommended_mode: "map",
161            },
162            RelevanceScore {
163                path: "/repo/website/src/a.astro".to_string(),
164                score: 0.05,
165                recommended_mode: "map",
166            },
167        ];
168        let refs: Vec<&RelevanceScore> = candidates.iter().collect();
169        let pop = decide_for_candidates(task, root, &refs);
170        assert!(pop.included_modules.contains(&"rust".to_string()));
171        assert!(pop.excluded_modules.iter().any(|e| e.module == "website"));
172        let kept = filter_candidates_by_pop(root, &refs, &pop);
173        assert!(kept.iter().all(|c| !c.path.contains("/website/")));
174    }
175
176    #[test]
177    fn pop_keeps_website_when_task_mentions_docs() {
178        let root = "/repo";
179        let task = "update website docs";
180        let candidates = [
181            RelevanceScore {
182                path: "/repo/website/src/index.astro".to_string(),
183                score: 0.2,
184                recommended_mode: "map",
185            },
186            RelevanceScore {
187                path: "/repo/rust/src/lib.rs".to_string(),
188                score: 0.2,
189                recommended_mode: "map",
190            },
191        ];
192        let refs: Vec<&RelevanceScore> = candidates.iter().collect();
193        let pop = decide_for_candidates(task, root, &refs);
194        assert!(pop.included_modules.contains(&"website".to_string()));
195        assert!(pop.excluded_modules.is_empty());
196    }
197}
lean_ctx/core/pop_pruning.rs

lean_ctx/core/
pop_pruning.rs
