Skip to main content

aft/compress/
generic.rs

1use crate::compress::{CompressionResult, Compressor};
2
3pub fn strip_ansi(input: &str) -> String {
4    record_ansi_strip();
5    let bytes = input.as_bytes();
6    let mut output = String::with_capacity(input.len());
7    let mut index = 0;
8    let mut last_kept = 0;
9
10    while index < bytes.len() {
11        if bytes[index] != 0x1b {
12            index += 1;
13            continue;
14        }
15
16        let Some(next) = bytes.get(index + 1).copied() else {
17            break;
18        };
19
20        let end = if next == b'[' {
21            let mut cursor = index + 2;
22            while cursor < bytes.len() {
23                if (0x40..=0x7e).contains(&bytes[cursor]) {
24                    cursor += 1;
25                    break;
26                }
27                cursor += 1;
28            }
29            cursor
30        } else if (0x40..=0x5f).contains(&next) {
31            index + 2
32        } else {
33            index += 1;
34            continue;
35        };
36
37        output.push_str(&input[last_kept..index]);
38        index = end.min(bytes.len());
39        last_kept = index;
40    }
41
42    output.push_str(&input[last_kept..]);
43    output
44}
45
46pub fn dedup_consecutive(input: &str) -> String {
47    let had_trailing_newline = input.ends_with('\n');
48    let mut output = String::with_capacity(input.len());
49    let mut lines = input.lines();
50
51    let Some(mut current) = lines.next() else {
52        return String::new();
53    };
54    let mut count = 1usize;
55
56    for line in lines {
57        if line == current {
58            count += 1;
59        } else {
60            push_dedup_run(&mut output, current, count);
61            current = line;
62            count = 1;
63        }
64    }
65    push_dedup_run(&mut output, current, count);
66
67    if !had_trailing_newline {
68        output.pop();
69    }
70
71    output
72}
73
74fn push_dedup_run(output: &mut String, line: &str, count: usize) {
75    output.push_str(line);
76    output.push('\n');
77    if count >= 4 {
78        output.push_str("... (");
79        output.push_str(&(count - 1).to_string());
80        output.push_str(" more)\n");
81    } else {
82        for _ in 1..count {
83            output.push_str(line);
84            output.push('\n');
85        }
86    }
87}
88
89pub fn middle_truncate(
90    input: &str,
91    threshold_bytes: usize,
92    keep_head: usize,
93    keep_tail: usize,
94) -> String {
95    if input.len() <= threshold_bytes {
96        return input.to_string();
97    }
98
99    let head_end = floor_char_boundary(input, keep_head.min(input.len()));
100    let tail_start = ceil_char_boundary(input, input.len().saturating_sub(keep_tail));
101
102    if head_end >= tail_start {
103        return input.to_string();
104    }
105
106    let truncated_bytes = tail_start - head_end;
107    let mut output = String::with_capacity(head_end + keep_tail + 64);
108    output.push_str(&input[..head_end]);
109    if !output.ends_with('\n') {
110        output.push('\n');
111    }
112    output.push_str("...<truncated ");
113    output.push_str(&truncated_bytes.to_string());
114    output.push_str(" bytes>...\n");
115    output.push_str(&input[tail_start..]);
116    output
117}
118
119pub(crate) fn floor_char_boundary(input: &str, mut index: usize) -> usize {
120    while index > 0 && !input.is_char_boundary(index) {
121        index -= 1;
122    }
123    index
124}
125
126pub(crate) fn ceil_char_boundary(input: &str, mut index: usize) -> usize {
127    while index < input.len() && !input.is_char_boundary(index) {
128        index += 1;
129    }
130    index
131}
132
133pub struct GenericCompressor;
134
135impl GenericCompressor {
136    pub fn compress_output(output: &str) -> String {
137        let stripped = strip_ansi(output);
138        Self::compress_stripped_output(&stripped)
139    }
140
141    pub(crate) fn compress_stripped_output(output: &str) -> String {
142        dedup_consecutive(output)
143    }
144}
145
146#[cfg(test)]
147thread_local! {
148    static ANSI_STRIP_COUNT: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
149}
150
151#[cfg(test)]
152fn record_ansi_strip() {
153    ANSI_STRIP_COUNT.with(|count| count.set(count.get() + 1));
154}
155
156#[cfg(not(test))]
157fn record_ansi_strip() {}
158
159#[cfg(test)]
160pub(crate) fn reset_ansi_strip_count() {
161    ANSI_STRIP_COUNT.with(|count| count.set(0));
162}
163
164#[cfg(test)]
165pub(crate) fn ansi_strip_count() -> usize {
166    ANSI_STRIP_COUNT.with(std::cell::Cell::get)
167}
168
169impl Compressor for GenericCompressor {
170    fn matches(&self, _command: &str) -> bool {
171        true
172    }
173
174    fn compress_with_exit_code(
175        &self,
176        _command: &str,
177        output: &str,
178        _exit_code: Option<i32>,
179    ) -> CompressionResult {
180        Self::compress_output(output).into()
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    #[test]
189    fn generic_does_not_pretruncate_above_old_five_kib_threshold() {
190        let input = (0..900)
191            .map(|idx| format!("unique-line-{idx:04}"))
192            .collect::<Vec<_>>()
193            .join("\n");
194        assert!(input.len() > 5 * 1024);
195
196        let compressed = GenericCompressor::compress_output(&input);
197
198        assert!(!compressed.contains("...<truncated "));
199        assert!(compressed.len() > 5 * 1024);
200        assert!(compressed.contains("unique-line-0000"));
201        assert!(compressed.contains("unique-line-0899"));
202    }
203}