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

lean_ctx/core/
symbol_map.rs
