lean_ctx/core/
symbol_map.rs

1use std::collections::HashMap;
2
3use crate::core::tokens::count_tokens;
4
5macro_rules! static_regex {
6    ($pattern:expr_2021) => {{
7        static RE: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
8        RE.get_or_init(|| {
9            regex::Regex::new($pattern).expect(concat!("BUG: invalid static regex: ", $pattern))
10        })
11    }};
12}
13
14const MIN_IDENT_LENGTH: usize = 6;
15const SHORT_ID_PREFIX: char = 'α';
16
17/// Whether alpha/§MAP identifier substitution should be applied to tool output.
18///
19/// Activation order:
20/// 1. `LEAN_CTX_SYMBOL_MAP=1` env var → force on
21/// 2. `LEAN_CTX_SYMBOL_MAP=0` env var → force off
22/// 3. `symbol_map_auto = true` in config + project >50 source files → auto-on
23/// 4. Default: off (the abbreviated form hinders editing; opt-in only)
24pub fn substitution_enabled() -> bool {
25    if let Ok(v) = std::env::var("LEAN_CTX_SYMBOL_MAP") {
26        return v == "1" || v.eq_ignore_ascii_case("true") || v.eq_ignore_ascii_case("on");
27    }
28    let cfg = crate::core::config::Config::load();
29    if cfg.symbol_map_auto {
30        return auto_detect_large_project();
31    }
32    false
33}
34
35fn auto_detect_large_project() -> bool {
36    use std::sync::OnceLock;
37    static DETECTED: OnceLock<bool> = OnceLock::new();
38    *DETECTED.get_or_init(|| {
39        let cwd = std::env::current_dir().unwrap_or_default();
40        let source_exts = [
41            "rs", "ts", "tsx", "js", "jsx", "py", "go", "java", "rb", "cpp", "c", "h", "gd",
42        ];
43        let count = ignore::WalkBuilder::new(&cwd)
44            .hidden(true)
45            .max_depth(Some(6))
46            .git_ignore(true)
47            .require_git(false)
48            .filter_entry(crate::core::walk_filter::keep_entry)
49            .build()
50            .filter_map(std::result::Result::ok)
51            .filter(|e| {
52                e.file_type().is_some_and(|ft| ft.is_file())
53                    && e.path()
54                        .extension()
55                        .and_then(|ext| ext.to_str())
56                        .is_some_and(|ext| source_exts.contains(&ext))
57            })
58            .take(51)
59            .count();
60        count > 50
61    })
62}
63
64#[derive(Debug, Clone)]
65pub struct SymbolMap {
66    forward: HashMap<String, String>,
67    next_id: usize,
68}
69
70impl Default for SymbolMap {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76impl SymbolMap {
77    pub fn new() -> Self {
78        Self {
79            forward: HashMap::new(),
80            next_id: 1,
81        }
82    }
83
84    pub fn register(&mut self, identifier: &str) -> Option<String> {
85        if identifier.len() < MIN_IDENT_LENGTH {
86            return None;
87        }
88
89        if let Some(existing) = self.forward.get(identifier) {
90            return Some(existing.clone());
91        }
92
93        let short_id = format!("{SHORT_ID_PREFIX}{}", self.next_id);
94        self.next_id += 1;
95        self.forward
96            .insert(identifier.to_string(), short_id.clone());
97        Some(short_id)
98    }
99
100    pub fn apply(&self, text: &str) -> String {
101        if self.forward.is_empty() {
102            return text.to_string();
103        }
104
105        let mut sorted: Vec<(&String, &String)> = self.forward.iter().collect();
106        sorted.sort_by_key(|x| std::cmp::Reverse(x.0.len()));
107
108        let mut result = text.to_string();
109        for (long, short) in &sorted {
110            result = result.replace(long.as_str(), short.as_str());
111        }
112        result
113    }
114
115    pub fn format_table(&self) -> String {
116        if self.forward.is_empty() {
117            return String::new();
118        }
119
120        let mut entries: Vec<(&String, &String)> = self.forward.iter().collect();
121        entries.sort_by_key(|(_, v)| {
122            v.trim_start_matches(SHORT_ID_PREFIX)
123                .parse::<usize>()
124                .unwrap_or(0)
125        });
126
127        let mut table = String::from("\n§MAP:");
128        for (long, short) in &entries {
129            table.push_str(&format!("\n  {short}={long}"));
130        }
131        table
132    }
133
134    pub fn len(&self) -> usize {
135        self.forward.len()
136    }
137
138    pub fn is_empty(&self) -> bool {
139        self.forward.is_empty()
140    }
141}
142
143/// MAP entry cost in tokens: "  αN=identifier\n" ≈ short_id_tokens + ident_tokens + 2 (= and newline)
144const MAP_ENTRY_OVERHEAD: usize = 2;
145
146/// ROI-based decision: register only when total savings exceed the MAP entry cost.
147/// savings = occurrences * (tokens(ident) - tokens(short_id))
148/// cost    = tokens(ident) + tokens(short_id) + MAP_ENTRY_OVERHEAD
149pub fn should_register(identifier: &str, occurrences: usize, next_id: usize) -> bool {
150    if identifier.len() < MIN_IDENT_LENGTH {
151        return false;
152    }
153    let ident_tokens = count_tokens(identifier);
154    let short_id = format!("{SHORT_ID_PREFIX}{next_id}");
155    let short_tokens = count_tokens(&short_id);
156
157    let token_saving_per_use = ident_tokens.saturating_sub(short_tokens);
158    if token_saving_per_use == 0 {
159        return false;
160    }
161
162    let total_savings = occurrences * token_saving_per_use;
163    let entry_cost = ident_tokens + short_tokens + MAP_ENTRY_OVERHEAD;
164
165    total_savings > entry_cost
166}
167
168pub fn extract_identifiers(content: &str, exts: &[&str]) -> Vec<String> {
169    let ident_re = static_regex!(r"\b[a-zA-Z_][a-zA-Z0-9_]*\b");
170
171    let mut seen = HashMap::new();
172    for mat in ident_re.find_iter(content) {
173        let word = mat.as_str();
174        if word.len() >= MIN_IDENT_LENGTH && !is_keyword(word, exts) {
175            *seen.entry(word.to_string()).or_insert(0usize) += 1;
176        }
177    }
178
179    let mut next_id = 1usize;
180    let mut idents: Vec<(String, usize)> = seen
181        .into_iter()
182        .filter(|(ident, count)| {
183            let pass = should_register(ident, *count, next_id);
184            if pass {
185                next_id += 1;
186            }
187            pass
188        })
189        .collect();
190
191    idents.sort_by(|a, b| {
192        let savings_a = a.0.len() * a.1;
193        let savings_b = b.0.len() * b.1;
194        savings_b.cmp(&savings_a)
195    });
196
197    idents.into_iter().map(|(s, _)| s).collect()
198}
199
200/// True when `word` is a language keyword for *any* of `exts`. An empty slice
201/// (no `include` glob, or one without a file extension) matches nothing, so all
202/// identifiers stay eligible for substitution.
203fn is_keyword(word: &str, exts: &[&str]) -> bool {
204    exts.iter().any(|&ext| match ext {
205        "rs" => matches!(
206            word,
207            "continue" | "default" | "return" | "struct" | "unsafe" | "where"
208        ),
209        "ts" | "tsx" | "js" | "jsx" => matches!(
210            word,
211            "constructor" | "arguments" | "undefined" | "prototype" | "instanceof"
212        ),
213        "py" => matches!(word, "continue" | "lambda" | "return" | "import" | "class"),
214        _ => false,
215    })
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    #[test]
223    fn test_should_register_short_ident_rejected() {
224        assert!(!should_register("foo", 100, 1));
225        assert!(!should_register("bar", 50, 1));
226        assert!(!should_register("x", 1000, 1));
227    }
228
229    #[test]
230    fn test_should_register_roi_positive() {
231        // Very long identifier (many BPE tokens) appearing 5 times
232        assert!(should_register(
233            "authenticate_user_credentials_handler",
234            5,
235            1
236        ));
237    }
238
239    #[test]
240    fn test_should_register_roi_negative_single_use() {
241        // Long ident but only 1 occurrence — MAP entry cost > savings
242        assert!(!should_register(
243            "authenticate_user_credentials_handler",
244            1,
245            1
246        ));
247    }
248
249    #[test]
250    fn test_should_register_roi_scales_with_frequency() {
251        let ident = "configuration_manager_instance";
252        // Should fail at low frequency, pass at high frequency
253        let passes_at_low = should_register(ident, 2, 1);
254        let passes_at_high = should_register(ident, 10, 1);
255        // At some point frequency makes it worthwhile
256        assert!(passes_at_high || !passes_at_low);
257    }
258
259    #[test]
260    fn test_extract_identifiers_roi_filtering() {
261        // Repeat a long identifier enough times that ROI is positive
262        let long = "authenticate_user_credentials_handler";
263        let content = format!("{long} {long} {long} {long} {long} short");
264        let result = extract_identifiers(&content, &["rs"]);
265        assert!(result.contains(&long.to_string()));
266        assert!(!result.contains(&"short".to_string()));
267    }
268
269    #[test]
270    fn test_register_returns_existing() {
271        let mut map = SymbolMap::new();
272        let first = map.register("validateToken");
273        let second = map.register("validateToken");
274        assert_eq!(first, second);
275    }
276
277    #[test]
278    fn test_apply_replaces_identifiers() {
279        let mut map = SymbolMap::new();
280        map.register("validateToken");
281        let result = map.apply("call validateToken here");
282        assert!(result.contains("α1"));
283        assert!(!result.contains("validateToken"));
284    }
285
286    #[test]
287    fn test_format_table_output() {
288        let mut map = SymbolMap::new();
289        map.register("validateToken");
290        let table = map.format_table();
291        assert!(table.contains("§MAP:"));
292        assert!(table.contains("α1=validateToken"));
293    }
294}
lean_ctx/core/symbol_map.rs

lean_ctx/core/
symbol_map.rs
