lean_ctx_patterns/
lossless_collapse.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Strictly lossless repeat collapse for the governor's R2 rule (governor
4//! design v1 §2 R2).
5//!
6//! A run of at least [`MIN_RUN`] byte-identical consecutive lines becomes the
7//! line once plus a marker line `[×N identical lines]`, where `N` counts the
8//! whole run. Nothing else changes: no timestamp stripping, no blank-line
9//! removal, no truncation, no reordering, no `\r` progress folding (folding
10//! progress rewrites drops bytes, so it is not lossless and is left out).
11//!
12//! [`expand`] inverts [`collapse`] exactly. To keep the marker unambiguous,
13//! [`collapse`] refuses input that already contains a marker-shaped line or the
14//! header line; the caller then keeps the original. This is deliberately not
15//! `log_dedup::compress`, which is lossy and stays as it is for `ctx_shell`.
16
17/// Shortest run of identical lines that is collapsed.
18pub const MIN_RUN: usize = 3;
19
20/// First line of a collapsed text, so the model knows the marker semantics.
21pub const HEADER_PREFIX: &str = "[lean-ctx: ";
22const HEADER_SUFFIX: &str = " repeated lines collapsed, lossless]";
23
24const MARKER_PREFIX: &str = "[×";
25const MARKER_SUFFIX: &str = " identical lines]";
26
27/// The collapse result: the rendered text and how many lines it removed.
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct Collapsed {
30    pub text: String,
31    pub lines_before: usize,
32    pub lines_after: usize,
33}
34
35fn marker_count(line: &str) -> Option<usize> {
36    line.strip_prefix(MARKER_PREFIX)?
37        .strip_suffix(MARKER_SUFFIX)?
38        .parse()
39        .ok()
40}
41
42fn is_reserved(line: &str) -> bool {
43    marker_count(line).is_some()
44        || (line.starts_with(HEADER_PREFIX) && line.ends_with(HEADER_SUFFIX))
45}
46
47/// Collapse runs of identical lines. `None` when nothing collapses or the input
48/// carries a reserved line (the result would not invert exactly).
49#[must_use]
50pub fn collapse(input: &str) -> Option<Collapsed> {
51    let lines: Vec<&str> = input.split('\n').collect();
52    if lines.iter().any(|l| is_reserved(l)) {
53        return None;
54    }
55    let mut out: Vec<String> = Vec::with_capacity(lines.len());
56    let mut removed = 0usize;
57    let mut i = 0;
58    while i < lines.len() {
59        let mut j = i + 1;
60        while j < lines.len() && lines[j] == lines[i] {
61            j += 1;
62        }
63        let run = j - i;
64        if run >= MIN_RUN {
65            out.push(lines[i].to_string());
66            out.push(format!("{MARKER_PREFIX}{run}{MARKER_SUFFIX}"));
67            removed += run - 2;
68        } else {
69            out.extend(lines[i..j].iter().map(|l| (*l).to_string()));
70        }
71        i = j;
72    }
73    if removed == 0 {
74        return None;
75    }
76    let body = out.join("\n");
77    Some(Collapsed {
78        text: format!("{HEADER_PREFIX}{removed}{HEADER_SUFFIX}\n{body}"),
79        lines_before: lines.len(),
80        lines_after: out.len() + 1,
81    })
82}
83
84/// Exact inverse of [`collapse`]. Text without the header is returned as is.
85#[must_use]
86pub fn expand(collapsed: &str) -> String {
87    let Some((header, body)) = collapsed.split_once('\n') else {
88        return collapsed.to_string();
89    };
90    if !(header.starts_with(HEADER_PREFIX) && header.ends_with(HEADER_SUFFIX)) {
91        return collapsed.to_string();
92    }
93    let mut out: Vec<&str> = Vec::new();
94    for line in body.split('\n') {
95        match (marker_count(line), out.last().copied()) {
96            (Some(n), Some(prev)) if n >= MIN_RUN => {
97                out.extend(std::iter::repeat_n(prev, n - 1));
98            }
99            _ => out.push(line),
100        }
101    }
102    out.join("\n")
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    /// Small deterministic generator: inputs built from a tiny alphabet so runs,
110    /// blank lines, trailing newlines and `\r` occur often.
111    fn generated_inputs() -> Vec<String> {
112        let atoms = [
113            "a",
114            "b",
115            "",
116            "Compiling x",
117            "\r50%\r100%",
118            "  indented",
119            "a ",
120        ];
121        let mut seed: u64 = 0x9E37_79B9_7F4A_7C15;
122        let mut out = Vec::new();
123        for _ in 0..2000 {
124            seed ^= seed << 13;
125            seed ^= seed >> 7;
126            seed ^= seed << 17;
127            let len = (seed % 40) as usize;
128            let mut lines = Vec::with_capacity(len);
129            let mut s = seed;
130            for _ in 0..len {
131                s = s.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
132                // Bias towards repeats of the previous line.
133                if !s.is_multiple_of(3) && !lines.is_empty() {
134                    let prev: &str = lines[lines.len() - 1];
135                    lines.push(prev);
136                } else {
137                    lines.push(atoms[(s >> 33) as usize % atoms.len()]);
138                }
139            }
140            let mut text = lines.join("\n");
141            if seed.is_multiple_of(2) {
142                text.push('\n');
143            }
144            out.push(text);
145        }
146        out
147    }
148
149    #[test]
150    fn expand_inverts_collapse_on_generated_inputs() {
151        let mut collapsed_any = 0;
152        for input in generated_inputs() {
153            if let Some(c) = collapse(&input) {
154                collapsed_any += 1;
155                assert_eq!(expand(&c.text), input, "round trip must be exact");
156                // Header line included; callers decide by measured savings.
157                assert!(c.lines_after <= c.lines_before);
158            }
159        }
160        assert!(collapsed_any > 100, "generator must exercise collapses");
161    }
162
163    #[test]
164    fn collapse_table() {
165        // (input, expected collapsed body or None)
166        let cases: &[(&str, Option<&str>)] = &[
167            ("a\nb\nc", None),
168            ("a\na", None),
169            ("a\na\na", Some("a\n[×3 identical lines]")),
170            (
171                "x\na\na\na\na\ny\n",
172                Some("x\na\n[×4 identical lines]\ny\n"),
173            ),
174            // Trailing blank lines are lines too; still exact.
175            ("a\n\n\n\n", Some("a\n\n[×4 identical lines]")),
176            // Near-identical lines (trailing space) never merge.
177            ("a\na \na\na ", None),
178            // A reserved marker in the input refuses the collapse.
179            ("[×3 identical lines]\nz\nz\nz", None),
180        ];
181        for (input, want) in cases {
182            let got = collapse(input).map(|c| {
183                let (_, body) = c.text.split_once('\n').unwrap();
184                body.to_string()
185            });
186            assert_eq!(got.as_deref(), *want, "input {input:?}");
187            if let Some(c) = collapse(input) {
188                assert_eq!(expand(&c.text), *input);
189            }
190        }
191    }
192
193    #[test]
194    fn header_counts_removed_lines() {
195        let c = collapse(&"same\n".repeat(10)).unwrap();
196        assert!(
197            c.text
198                .starts_with("[lean-ctx: 8 repeated lines collapsed, lossless]\n")
199        );
200    }
201
202    #[test]
203    fn expand_leaves_plain_text_alone() {
204        for t in ["", "plain", "a\n[×3 identical lines]"] {
205            assert_eq!(expand(t), t);
206        }
207    }
208}
lean_ctx_patterns/lossless_collapse.rs

lean_ctx_patterns/
lossless_collapse.rs
