Skip to main content

cersei_compression/
truncate.rs

1//! Line-structure-aware truncation — preserves function signatures and
2//! imports, omits the body.
3//!
4//! Credits: adapted from rtk (Rust Token Killer) — `rtk/src/core/filter.rs`.
5//! MIT © Patrick Szymkowiak. See LICENSE.
6
7use once_cell::sync::Lazy;
8use regex::Regex;
9
10static IMPORT_PATTERN: Lazy<Regex> =
11    Lazy::new(|| Regex::new(r"^(use |import |from |require\(|#include)").unwrap());
12static FUNC_SIGNATURE: Lazy<Regex> = Lazy::new(|| {
13    Regex::new(
14        r"^(pub\s+)?(async\s+)?(fn|def|function|func|class|struct|enum|trait|interface|type)\s+\w+",
15    )
16    .unwrap()
17});
18
19pub fn smart_truncate(content: &str, max_lines: usize) -> String {
20    let lines: Vec<&str> = content.lines().collect();
21    if lines.len() <= max_lines {
22        return content.to_string();
23    }
24
25    let mut out: Vec<String> = Vec::with_capacity(max_lines);
26    let mut kept = 0usize;
27    let mut skipped = false;
28
29    for line in &lines {
30        let trimmed = line.trim();
31        let important = FUNC_SIGNATURE.is_match(trimmed)
32            || IMPORT_PATTERN.is_match(trimmed)
33            || trimmed.starts_with("pub ")
34            || trimmed.starts_with("export ")
35            || trimmed == "}"
36            || trimmed == "{";
37
38        if important || kept < max_lines / 2 {
39            if skipped {
40                out.push(format!(
41                    "    // ... {} lines omitted",
42                    lines.len().saturating_sub(kept)
43                ));
44                skipped = false;
45            }
46            out.push((*line).to_string());
47            kept += 1;
48        } else {
49            skipped = true;
50        }
51
52        if kept >= max_lines.saturating_sub(1) {
53            break;
54        }
55    }
56
57    if skipped || kept < lines.len() {
58        out.push(format!(
59            "// ... {} more lines (total: {})",
60            lines.len().saturating_sub(kept),
61            lines.len()
62        ));
63    }
64
65    out.join("\n")
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn noop_when_short() {
74        let src = "a\nb\nc";
75        assert_eq!(smart_truncate(src, 100), src);
76    }
77
78    #[test]
79    fn truncates_and_counts() {
80        let src: String = (0..200)
81            .map(|i| format!("plain text line {i}"))
82            .collect::<Vec<_>>()
83            .join("\n");
84        let out = smart_truncate(&src, 20);
85        let overflow = out
86            .lines()
87            .find(|l| l.contains("more lines"))
88            .expect("overflow line");
89        let n: usize = overflow
90            .split_whitespace()
91            .find_map(|w| w.parse().ok())
92            .expect("count");
93        let kept = out
94            .lines()
95            .filter(|l| !l.contains("more lines") && !l.contains("omitted"))
96            .count();
97        assert_eq!(kept + n, 200);
98    }
99}