lean_ctx/core/code_health/
annotate.rs

1//! Read-time code-health annotations.
2//!
3//! Produces **sparse, deterministic** per-function annotations (`cc=18`) for the
4//! over-threshold functions in a file, so an agent reading signatures or the map
5//! sees complexity hotspots inline and can decide *not* to read a giant function
6//! in full. Only functions above the threshold are annotated, keeping output
7//! byte-stable and cheap (#498-safe).
8
9use super::cognitive::cognitive_per_function;
10use std::collections::HashMap;
11
12/// One inline annotation for a function in a read view.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct ReadAnnotation {
15    /// 1-based start line of the function.
16    pub line: usize,
17    pub name: String,
18    /// The compact note to append, e.g. `cc=18`.
19    pub note: String,
20}
21
22/// Annotations for the over-threshold functions in `source` of file `ext`.
23/// Sorted by line then name. Empty when tree-sitter is off, the extension is
24/// unsupported, or nothing is over threshold.
25pub fn annotations_for_file(source: &str, ext: &str, threshold: u32) -> Vec<ReadAnnotation> {
26    let mut out: Vec<ReadAnnotation> = Vec::new();
27    if let Some(fns) = cognitive_per_function(source, ext) {
28        for f in fns {
29            if f.cognitive > threshold {
30                out.push(ReadAnnotation {
31                    line: f.line,
32                    name: f.name,
33                    note: format!("cc={}", f.cognitive),
34                });
35            }
36        }
37    }
38    out.sort_by(|a, b| a.line.cmp(&b.line).then_with(|| a.name.cmp(&b.name)));
39    out
40}
41
42/// Index annotations by function name for renderers that match on the symbol
43/// name (more robust than line numbers when attributes/decorators shift lines).
44pub fn by_name(annotations: &[ReadAnnotation]) -> HashMap<String, String> {
45    annotations
46        .iter()
47        .map(|a| (a.name.clone(), a.note.clone()))
48        .collect()
49}
50
51/// Cognitive complexity of the function named `name` defined nearest `start_line`
52/// in `source` (file extension `ext`). Unlike [`annotations_for_file`] this is
53/// *not* threshold-gated, so a targeted `ctx_symbol` query can show the exact cc
54/// of any function. `None` when tree-sitter is off, the extension is
55/// unsupported, or no function with that name is present (e.g. a struct symbol).
56pub fn cognitive_for_symbol(source: &str, ext: &str, name: &str, start_line: usize) -> Option<u32> {
57    let fns = cognitive_per_function(source, ext)?;
58    fns.iter()
59        .filter(|f| f.name == name)
60        .min_by_key(|f| f.line.abs_diff(start_line))
61        .map(|f| f.cognitive)
62}
63
64#[cfg(all(test, feature = "tree-sitter"))]
65pub mod tests {
66    use super::*;
67
68    #[test]
69    fn annotates_only_over_threshold() {
70        // `deep` = 1+2+3+4 = 10 cognitive; `flat` = 0.
71        let src = "fn flat() {}\nfn deep(a: bool) { if a { if a { if a { if a {} } } } }\n";
72        let anns = annotations_for_file(src, "rs", 5);
73        assert_eq!(anns.len(), 1);
74        assert_eq!(anns[0].name, "deep");
75        assert_eq!(anns[0].note, "cc=10");
76    }
77
78    #[test]
79    fn nothing_when_under_threshold() {
80        let src = "fn small(a: bool) { if a {} }\n";
81        {
82            let found = &annotations_for_file(src, "rs", 15);
83            assert!(found.is_empty(), "{found:?}");
84        }
85    }
86
87    #[test]
88    fn deterministic_across_runs() {
89        let src = "fn deep(a: bool) { if a { if a { if a { if a {} } } } }\n";
90        assert_eq!(
91            annotations_for_file(src, "rs", 5),
92            annotations_for_file(src, "rs", 5)
93        );
94    }
95
96    #[test]
97    fn by_name_lookup() {
98        let src = "fn deep(a: bool) { if a { if a { if a { if a {} } } } }\n";
99        let anns = annotations_for_file(src, "rs", 5);
100        let map = by_name(&anns);
101        assert_eq!(map.get("deep").map(String::as_str), Some("cc=10"));
102    }
103
104    #[test]
105    fn cognitive_for_symbol_reports_any_function() {
106        let src = "fn flat() {}\nfn deep(a: bool) { if a { if a { if a {} } } }\n";
107        // Not threshold-gated: even the flat function resolves (cc=0).
108        assert_eq!(cognitive_for_symbol(src, "rs", "flat", 1), Some(0));
109        assert_eq!(cognitive_for_symbol(src, "rs", "deep", 2), Some(6));
110        assert_eq!(cognitive_for_symbol(src, "rs", "missing", 1), None);
111    }
112
113    #[test]
114    fn cognitive_for_symbol_disambiguates_by_line() {
115        // Two same-named functions: pick the one nearest the queried line.
116        let src = "fn dup(a: bool) { if a {} }\nfn other() {}\nfn dup(a: bool) { if a { if a { if a {} } } }\n";
117        assert_eq!(cognitive_for_symbol(src, "rs", "dup", 1), Some(1));
118        assert_eq!(cognitive_for_symbol(src, "rs", "dup", 3), Some(6));
119    }
120}
lean_ctx/core/code_health/annotate.rs

lean_ctx/core/code_health/
annotate.rs
