kedge_compact/
lib.rs

1//! # kedge-compact
2//!
3//! Context-window management by *AST-aware* compaction. When a source file is
4//! too large to fit the budget, we don't truncate blindly — we parse it with
5//! Tree-sitter and keep the structural skeleton (every function/method/type
6//! signature, doc comments, imports) while eliding function *bodies*, which are
7//! usually the bulk of the tokens and the least useful for high-level reasoning.
8//!
9//! Five languages are supported — **Rust, Python, JavaScript, TypeScript, Go** —
10//! selected explicitly or by file extension. Each knows which node kinds are
11//! "functions" and how to spell an elided body in its own syntax.
12
13use std::path::Path;
14
15use serde::Serialize;
16use thiserror::Error;
17use tree_sitter::{Node, Parser};
18
19use kedge_core::HarnessError;
20
21#[derive(Debug, Error)]
22pub enum CompactError {
23    #[error("failed to load Tree-sitter grammar: {0}")]
24    Language(#[from] tree_sitter::LanguageError),
25    #[error("parser produced no tree")]
26    ParseFailed,
27    #[error("unsupported language for `.{0}` (supported: rs, py, js, ts, go)")]
28    UnsupportedExtension(String),
29}
30
31impl From<CompactError> for HarnessError {
32    fn from(e: CompactError) -> Self {
33        HarnessError::backend("compact", e)
34    }
35}
36
37pub type Result<T> = std::result::Result<T, CompactError>;
38
39/// A source language the compactor understands.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum Language {
42    Rust,
43    Python,
44    JavaScript,
45    TypeScript,
46    Go,
47}
48
49impl Language {
50    pub fn name(self) -> &'static str {
51        match self {
52            Language::Rust => "rust",
53            Language::Python => "python",
54            Language::JavaScript => "javascript",
55            Language::TypeScript => "typescript",
56            Language::Go => "go",
57        }
58    }
59
60    /// Map a file extension (without the dot) to a language.
61    pub fn from_extension(ext: &str) -> Option<Self> {
62        Some(match ext.to_ascii_lowercase().as_str() {
63            "rs" => Language::Rust,
64            "py" | "pyi" => Language::Python,
65            "js" | "jsx" | "mjs" | "cjs" => Language::JavaScript,
66            "ts" | "mts" | "cts" => Language::TypeScript,
67            "go" => Language::Go,
68            _ => return None,
69        })
70    }
71
72    /// Detect a language from a path's extension.
73    pub fn from_path(path: impl AsRef<Path>) -> Option<Self> {
74        path.as_ref()
75            .extension()
76            .and_then(|e| e.to_str())
77            .and_then(Self::from_extension)
78    }
79
80    fn grammar(self) -> tree_sitter::Language {
81        match self {
82            Language::Rust => tree_sitter_rust::LANGUAGE.into(),
83            Language::Python => tree_sitter_python::LANGUAGE.into(),
84            Language::JavaScript => tree_sitter_javascript::LANGUAGE.into(),
85            Language::TypeScript => tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
86            Language::Go => tree_sitter_go::LANGUAGE.into(),
87        }
88    }
89
90    /// Node kinds whose `body` field we elide.
91    fn function_kinds(self) -> &'static [&'static str] {
92        match self {
93            Language::Rust => &["function_item"],
94            Language::Python => &["function_definition"],
95            Language::JavaScript | Language::TypeScript => &[
96                "function_declaration",
97                "generator_function_declaration",
98                "method_definition",
99                "function_expression",
100            ],
101            Language::Go => &["function_declaration", "method_declaration"],
102        }
103    }
104
105    /// How to spell an elided body in this language's syntax.
106    fn body_placeholder(self, lines: usize) -> String {
107        match self {
108            // Python bodies are indented suites, not brace blocks.
109            Language::Python => format!("...  # {lines} lines elided"),
110            _ => format!("{{ /* {lines} lines elided */ }}"),
111        }
112    }
113}
114
115/// Approximate token count.
116///
117/// A real BPE tokenizer is model-specific; for budgeting we use the widely-used
118/// heuristic of ~3.7 characters per token, which tracks code closely enough to
119/// drive compaction decisions.
120pub fn estimate_tokens(text: &str) -> usize {
121    const CHARS_PER_TOKEN: f64 = 3.7;
122    (text.chars().count() as f64 / CHARS_PER_TOKEN).ceil() as usize
123}
124
125/// The outcome of a compaction pass.
126#[derive(Debug, Clone, Serialize)]
127pub struct CompactResult {
128    pub text: String,
129    pub original_tokens: usize,
130    pub compacted_tokens: usize,
131    /// How many function bodies were elided.
132    pub elided_bodies: usize,
133}
134
135impl CompactResult {
136    /// Fraction of tokens removed (0.0 = unchanged).
137    pub fn savings(&self) -> f64 {
138        if self.original_tokens == 0 {
139            return 0.0;
140        }
141        1.0 - (self.compacted_tokens as f64 / self.original_tokens as f64)
142    }
143}
144
145/// AST-aware compactor bound to one language.
146pub struct Compactor {
147    parser: Parser,
148    language: Language,
149}
150
151impl Compactor {
152    /// Create a compactor for a given language.
153    pub fn new(language: Language) -> Result<Self> {
154        let mut parser = Parser::new();
155        parser.set_language(&language.grammar())?;
156        Ok(Compactor { parser, language })
157    }
158
159    /// Convenience for the Rust grammar.
160    pub fn rust() -> Result<Self> {
161        Self::new(Language::Rust)
162    }
163
164    /// Build a compactor by detecting the language from a path's extension.
165    pub fn for_path(path: impl AsRef<Path>) -> Result<Self> {
166        let path = path.as_ref();
167        let lang = Language::from_path(path).ok_or_else(|| {
168            CompactError::UnsupportedExtension(
169                path.extension()
170                    .and_then(|e| e.to_str())
171                    .unwrap_or("")
172                    .to_string(),
173            )
174        })?;
175        Self::new(lang)
176    }
177
178    pub fn language(&self) -> Language {
179        self.language
180    }
181
182    /// Produce a structural outline: every signature kept, every function body
183    /// replaced with a one-line placeholder recording how much was elided.
184    #[tracing::instrument(
185        name = "compact.outline",
186        skip_all,
187        fields(
188            language = ?self.language,
189            files_compacted = 1,
190            tokens_saved = tracing::field::Empty,
191            elided_bodies = tracing::field::Empty,
192        )
193    )]
194    pub fn outline(&mut self, source: &str) -> Result<CompactResult> {
195        let tree = self
196            .parser
197            .parse(source, None)
198            .ok_or(CompactError::ParseFailed)?;
199        let mut edits: Vec<(usize, usize, String)> = Vec::new();
200        collect_body_elisions(tree.root_node(), source, self.language, &mut edits, 0);
201        edits.sort_by_key(|e| e.0);
202
203        let mut out = String::with_capacity(source.len());
204        let mut pos = 0;
205        for (start, end, replacement) in &edits {
206            if *start < pos {
207                continue; // guard against any overlap
208            }
209            out.push_str(&source[pos..*start]);
210            out.push_str(replacement);
211            pos = *end;
212        }
213        out.push_str(&source[pos..]);
214
215        let result = CompactResult {
216            original_tokens: estimate_tokens(source),
217            compacted_tokens: estimate_tokens(&out),
218            elided_bodies: edits.len(),
219            text: out,
220        };
221        let span = tracing::Span::current();
222        span.record(
223            "tokens_saved",
224            result
225                .original_tokens
226                .saturating_sub(result.compacted_tokens) as u64,
227        );
228        span.record("elided_bodies", result.elided_bodies as u64);
229        Ok(result)
230    }
231
232    /// Return `source` untouched if it already fits `max_tokens`; otherwise
233    /// compact it to an outline (best-effort — a file that is all signatures may
234    /// still exceed the budget).
235    pub fn compact_to_budget(&mut self, source: &str, max_tokens: usize) -> Result<CompactResult> {
236        let original = estimate_tokens(source);
237        if original <= max_tokens {
238            return Ok(CompactResult {
239                text: source.to_string(),
240                original_tokens: original,
241                compacted_tokens: original,
242                elided_bodies: 0,
243            });
244        }
245        self.outline(source)
246    }
247
248    /// The inverse of compaction: return the full source of every function or
249    /// method named `symbol`, bodies included.
250    ///
251    /// This is what makes compaction safe to *act* on. A caller orients on a
252    /// skeleton (bodies elided, signatures kept), then — when it needs to edit
253    /// one specific function whose body it never saw — fetches just that body on
254    /// demand instead of re-reading the whole file. The skeleton already shows
255    /// the name above each `/* … elided */`, so the caller always knows what to
256    /// ask for.
257    ///
258    /// Names are not unique (several `new` methods, an overridden `render`), so
259    /// every match is returned in source order. An empty result means no such
260    /// named function exists — the caller should fall back to reading the file.
261    #[tracing::instrument(
262        name = "compact.expand",
263        skip_all,
264        fields(language = ?self.language, symbol = %symbol, matches = tracing::field::Empty)
265    )]
266    pub fn expand(&mut self, source: &str, symbol: &str) -> Result<Vec<String>> {
267        let tree = self
268            .parser
269            .parse(source, None)
270            .ok_or(CompactError::ParseFailed)?;
271        let mut hits = Vec::new();
272        collect_named(
273            tree.root_node(),
274            source,
275            self.language,
276            symbol,
277            &mut hits,
278            0,
279        );
280        tracing::Span::current().record("matches", hits.len() as u64);
281        Ok(hits)
282    }
283}
284
285/// Cap on AST recursion depth. Native stacks handle far more, but a pathologically
286/// nested source (thousands of nested blocks) could otherwise overflow the stack —
287/// and a Rust stack overflow aborts the process, uncatchable. Real code never
288/// nests this deep; beyond the cap we simply stop descending (a few function
289/// bodies that deep may go un-elided — a cosmetic loss, never a crash).
290const MAX_AST_DEPTH: u32 = 400;
291
292/// Walk the tree collecting the full text of every function-kind node whose
293/// `name` field equals `symbol`. Mirrors [`collect_body_elisions`] but keeps the
294/// body instead of eliding it, and does not recurse into a matched function.
295fn collect_named(
296    node: Node,
297    source: &str,
298    lang: Language,
299    symbol: &str,
300    out: &mut Vec<String>,
301    depth: u32,
302) {
303    if depth > MAX_AST_DEPTH {
304        return;
305    }
306    let kinds = lang.function_kinds();
307    let mut cursor = node.walk();
308    for child in node.children(&mut cursor) {
309        if kinds.contains(&child.kind()) {
310            let name = child
311                .child_by_field_name("name")
312                .and_then(|n| n.utf8_text(source.as_bytes()).ok());
313            if name == Some(symbol) {
314                if let Ok(text) = child.utf8_text(source.as_bytes()) {
315                    out.push(text.to_string());
316                }
317            }
318            // A match's body is what we want; a non-match function can't contain
319            // a top-level-named target either — in both cases, don't recurse.
320        } else {
321            collect_named(child, source, lang, symbol, out, depth + 1);
322        }
323    }
324}
325
326/// Walk the tree collecting `(start, end, replacement)` for each function body,
327/// without descending into a body once found (its contents are elided anyway).
328fn collect_body_elisions(
329    node: Node,
330    source: &str,
331    lang: Language,
332    edits: &mut Vec<(usize, usize, String)>,
333    depth: u32,
334) {
335    if depth > MAX_AST_DEPTH {
336        return;
337    }
338    let kinds = lang.function_kinds();
339    let mut cursor = node.walk();
340    for child in node.children(&mut cursor) {
341        if kinds.contains(&child.kind()) {
342            if let Some(body) = child.child_by_field_name("body") {
343                let span = &source[body.start_byte()..body.end_byte()];
344                let lines = span.lines().count().max(1);
345                let replacement = lang.body_placeholder(lines);
346                // Only elide when it actually shrinks the source.
347                if replacement.len() < span.len() {
348                    edits.push((body.start_byte(), body.end_byte(), replacement));
349                }
350            }
351            // Do not recurse into the function; its body is gone.
352        } else {
353            // Descend into impls, classes, modules, etc. to reach nested functions.
354            collect_body_elisions(child, source, lang, edits, depth + 1);
355        }
356    }
357}
358
359#[cfg(test)]
360mod tests {
361    use super::*;
362
363    const RUST: &str = r#"
364/// A widget.
365pub struct Widget { pub id: u32 }
366impl Widget {
367    /// Build one.
368    pub fn new(id: u32) -> Self {
369        let secret = compute_expensive_thing(id);
370        Widget { id: secret }
371    }
372}
373pub fn top_level(a: i32, b: i32) -> i32 {
374    let mut acc = 0;
375    for i in a..b { acc += i * i; }
376    acc
377}
378"#;
379
380    #[test]
381    fn estimate_is_monotonic() {
382        assert!(estimate_tokens("short") < estimate_tokens("a considerably longer string here"));
383        assert_eq!(estimate_tokens(""), 0);
384    }
385
386    #[test]
387    fn rust_outline_keeps_signatures_and_elides_bodies() {
388        let mut c = Compactor::rust().unwrap();
389        let r = c.outline(RUST).unwrap();
390        assert!(r.text.contains("pub struct Widget"));
391        assert!(r.text.contains("pub fn new(id: u32) -> Self"));
392        assert!(r.text.contains("pub fn top_level(a: i32, b: i32) -> i32"));
393        assert!(r.text.contains("/// Build one."));
394        assert!(!r.text.contains("compute_expensive_thing"));
395        assert!(!r.text.contains("acc += i * i"));
396        assert_eq!(r.elided_bodies, 2);
397        assert!(r.compacted_tokens < r.original_tokens);
398    }
399
400    #[test]
401    fn python_outline_keeps_defs_and_uses_ellipsis() {
402        let src = "def compute(a, b):\n    total = 0\n    for i in range(a, b):\n        total += i * i\n    return total\n\nclass Widget:\n    def build(self, n):\n        secret = expensive(n)\n        return secret\n";
403        let mut c = Compactor::new(Language::Python).unwrap();
404        let r = c.outline(src).unwrap();
405        assert!(r.text.contains("def compute(a, b):"));
406        assert!(r.text.contains("def build(self, n):"));
407        assert!(r.text.contains("class Widget:"));
408        assert!(r.text.contains("..."));
409        assert!(!r.text.contains("total += i * i"));
410        assert!(!r.text.contains("expensive(n)"));
411        assert_eq!(r.elided_bodies, 2);
412    }
413
414    #[test]
415    fn javascript_outline_elides_function_and_method_bodies() {
416        let src = "function add(a, b) {\n  const s = a + b;\n  return s;\n}\nclass C {\n  method(x) {\n    const y = heavyThing(x);\n    return y;\n  }\n}\n";
417        let mut c = Compactor::new(Language::JavaScript).unwrap();
418        let r = c.outline(src).unwrap();
419        assert!(r.text.contains("function add(a, b)"));
420        assert!(r.text.contains("method(x)"));
421        assert!(!r.text.contains("heavyThing"));
422        assert_eq!(r.elided_bodies, 2);
423    }
424
425    #[test]
426    fn go_outline_elides_func_and_method_bodies() {
427        let src = "package main\nfunc Add(a, b int) int {\n\ts := a + b\n\treturn s\n}\nfunc (w Widget) Build(n int) int {\n\tsecret := expensive(n)\n\treturn secret\n}\n";
428        let mut c = Compactor::new(Language::Go).unwrap();
429        let r = c.outline(src).unwrap();
430        assert!(r.text.contains("func Add(a, b int) int"));
431        assert!(r.text.contains("func (w Widget) Build(n int) int"));
432        assert!(!r.text.contains("expensive(n)"));
433        assert_eq!(r.elided_bodies, 2);
434    }
435
436    #[test]
437    fn language_detection_by_extension() {
438        assert_eq!(Language::from_extension("rs"), Some(Language::Rust));
439        assert_eq!(Language::from_extension("PY"), Some(Language::Python));
440        assert_eq!(Language::from_extension("tsx"), None); // tsx not wired
441        assert_eq!(Language::from_extension("ts"), Some(Language::TypeScript));
442        assert_eq!(Language::from_extension("go"), Some(Language::Go));
443        assert_eq!(Language::from_extension("txt"), None);
444        assert_eq!(Language::from_path("src/main.rs"), Some(Language::Rust));
445    }
446
447    #[test]
448    fn for_path_rejects_unknown_extensions() {
449        assert!(matches!(
450            Compactor::for_path("data.txt"),
451            Err(CompactError::UnsupportedExtension(_))
452        ));
453    }
454
455    #[test]
456    fn compaction_never_grows_the_source() {
457        let mut c = Compactor::rust().unwrap();
458        let tiny = "fn a() { 1 }\nfn b() -> i32 { 2 }\n";
459        let r = c.outline(tiny).unwrap();
460        assert!(r.compacted_tokens <= r.original_tokens);
461        assert_eq!(r.elided_bodies, 0);
462    }
463
464    // ── expand: the inverse of compaction (JIT decompaction) ──
465
466    #[test]
467    fn expand_returns_the_full_body_a_skeleton_elided() {
468        let mut c = Compactor::rust().unwrap();
469        // The skeleton would show `pub fn new(...)` with the body elided; expand
470        // brings back exactly the code you'd need to edit it.
471        let hits = c.expand(RUST, "new").unwrap();
472        assert_eq!(hits.len(), 1);
473        assert!(hits[0].contains("compute_expensive_thing(id)"));
474        assert!(hits[0].contains("pub fn new"));
475    }
476
477    #[test]
478    fn expand_finds_top_level_functions_too() {
479        let mut c = Compactor::rust().unwrap();
480        let hits = c.expand(RUST, "top_level").unwrap();
481        assert_eq!(hits.len(), 1);
482        assert!(hits[0].contains("acc += i * i"));
483    }
484
485    #[test]
486    fn expand_returns_every_match_since_names_are_not_unique() {
487        let mut c = Compactor::rust().unwrap();
488        let src = "impl A { fn make() -> u8 { 1 } }\nimpl B { fn make() -> u8 { 2 } }\n";
489        let hits = c.expand(src, "make").unwrap();
490        assert_eq!(hits.len(), 2, "both `make` methods should come back");
491    }
492
493    #[test]
494    fn expand_of_an_unknown_symbol_is_empty_not_an_error() {
495        let mut c = Compactor::rust().unwrap();
496        assert!(c.expand(RUST, "does_not_exist").unwrap().is_empty());
497    }
498
499    #[test]
500    fn expand_works_across_languages() {
501        let mut py = Compactor::new(Language::Python).unwrap();
502        let hits = py
503            .expand("def greet(name):\n    return f'hi {name}'\n", "greet")
504            .unwrap();
505        assert_eq!(hits.len(), 1);
506        assert!(hits[0].contains("return f'hi {name}'"));
507    }
508}
kedge_compact/lib.rs

kedge_compact/
lib.rs
