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);
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
249/// Walk the tree collecting `(start, end, replacement)` for each function body,
250/// without descending into a body once found (its contents are elided anyway).
251fn collect_body_elisions(
252    node: Node,
253    source: &str,
254    lang: Language,
255    edits: &mut Vec<(usize, usize, String)>,
256) {
257    let kinds = lang.function_kinds();
258    let mut cursor = node.walk();
259    for child in node.children(&mut cursor) {
260        if kinds.contains(&child.kind()) {
261            if let Some(body) = child.child_by_field_name("body") {
262                let span = &source[body.start_byte()..body.end_byte()];
263                let lines = span.lines().count().max(1);
264                let replacement = lang.body_placeholder(lines);
265                // Only elide when it actually shrinks the source.
266                if replacement.len() < span.len() {
267                    edits.push((body.start_byte(), body.end_byte(), replacement));
268                }
269            }
270            // Do not recurse into the function; its body is gone.
271        } else {
272            // Descend into impls, classes, modules, etc. to reach nested functions.
273            collect_body_elisions(child, source, lang, edits);
274        }
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    const RUST: &str = r#"
283/// A widget.
284pub struct Widget { pub id: u32 }
285impl Widget {
286    /// Build one.
287    pub fn new(id: u32) -> Self {
288        let secret = compute_expensive_thing(id);
289        Widget { id: secret }
290    }
291}
292pub fn top_level(a: i32, b: i32) -> i32 {
293    let mut acc = 0;
294    for i in a..b { acc += i * i; }
295    acc
296}
297"#;
298
299    #[test]
300    fn estimate_is_monotonic() {
301        assert!(estimate_tokens("short") < estimate_tokens("a considerably longer string here"));
302        assert_eq!(estimate_tokens(""), 0);
303    }
304
305    #[test]
306    fn rust_outline_keeps_signatures_and_elides_bodies() {
307        let mut c = Compactor::rust().unwrap();
308        let r = c.outline(RUST).unwrap();
309        assert!(r.text.contains("pub struct Widget"));
310        assert!(r.text.contains("pub fn new(id: u32) -> Self"));
311        assert!(r.text.contains("pub fn top_level(a: i32, b: i32) -> i32"));
312        assert!(r.text.contains("/// Build one."));
313        assert!(!r.text.contains("compute_expensive_thing"));
314        assert!(!r.text.contains("acc += i * i"));
315        assert_eq!(r.elided_bodies, 2);
316        assert!(r.compacted_tokens < r.original_tokens);
317    }
318
319    #[test]
320    fn python_outline_keeps_defs_and_uses_ellipsis() {
321        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";
322        let mut c = Compactor::new(Language::Python).unwrap();
323        let r = c.outline(src).unwrap();
324        assert!(r.text.contains("def compute(a, b):"));
325        assert!(r.text.contains("def build(self, n):"));
326        assert!(r.text.contains("class Widget:"));
327        assert!(r.text.contains("..."));
328        assert!(!r.text.contains("total += i * i"));
329        assert!(!r.text.contains("expensive(n)"));
330        assert_eq!(r.elided_bodies, 2);
331    }
332
333    #[test]
334    fn javascript_outline_elides_function_and_method_bodies() {
335        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";
336        let mut c = Compactor::new(Language::JavaScript).unwrap();
337        let r = c.outline(src).unwrap();
338        assert!(r.text.contains("function add(a, b)"));
339        assert!(r.text.contains("method(x)"));
340        assert!(!r.text.contains("heavyThing"));
341        assert_eq!(r.elided_bodies, 2);
342    }
343
344    #[test]
345    fn go_outline_elides_func_and_method_bodies() {
346        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";
347        let mut c = Compactor::new(Language::Go).unwrap();
348        let r = c.outline(src).unwrap();
349        assert!(r.text.contains("func Add(a, b int) int"));
350        assert!(r.text.contains("func (w Widget) Build(n int) int"));
351        assert!(!r.text.contains("expensive(n)"));
352        assert_eq!(r.elided_bodies, 2);
353    }
354
355    #[test]
356    fn language_detection_by_extension() {
357        assert_eq!(Language::from_extension("rs"), Some(Language::Rust));
358        assert_eq!(Language::from_extension("PY"), Some(Language::Python));
359        assert_eq!(Language::from_extension("tsx"), None); // tsx not wired
360        assert_eq!(Language::from_extension("ts"), Some(Language::TypeScript));
361        assert_eq!(Language::from_extension("go"), Some(Language::Go));
362        assert_eq!(Language::from_extension("txt"), None);
363        assert_eq!(Language::from_path("src/main.rs"), Some(Language::Rust));
364    }
365
366    #[test]
367    fn for_path_rejects_unknown_extensions() {
368        assert!(matches!(
369            Compactor::for_path("data.txt"),
370            Err(CompactError::UnsupportedExtension(_))
371        ));
372    }
373
374    #[test]
375    fn compaction_never_grows_the_source() {
376        let mut c = Compactor::rust().unwrap();
377        let tiny = "fn a() { 1 }\nfn b() -> i32 { 2 }\n";
378        let r = c.outline(tiny).unwrap();
379        assert!(r.compacted_tokens <= r.original_tokens);
380        assert_eq!(r.elided_bodies, 0);
381    }
382}
kedge_compact/lib.rs

kedge_compact/
lib.rs
