1pub trait SliceExt {
3 fn trim(&self) -> &Self;
4 fn trim_trailing(&self) -> &Self;
5}
6
7fn is_whitespace(c: &u8) -> bool {
8 *c == b'\t' || *c == b' '
9}
10
11fn is_not_whitespace(c: &u8) -> bool {
12 !is_whitespace(c)
13}
14
15impl SliceExt for [u8] {
16 fn trim(&self) -> &[u8] {
18 if let Some(first) = self.iter().position(is_not_whitespace) {
19 if let Some(last) = self.iter().rposition(is_not_whitespace) {
20 &self[first..last + 1]
21 } else {
22 unreachable!();
23 }
24 } else {
25 &[]
26 }
27 }
28
29 fn trim_trailing(&self) -> &[u8] {
31 if let Some(last) = self.iter().rposition(is_not_whitespace) {
32 &self[0..last + 1]
33 } else {
34 &[]
35 }
36 }
37}
38
39fn buf_trim_trailing(buf: &[u8]) -> Vec<&[u8]> {
41 let trimmed_lines: Vec<&[u8]> = buf
42 .split(|c| *c == b'\n')
43 .map(SliceExt::trim_trailing) .collect();
45
46 trimmed_lines
47}
48
49fn drop_last_empty_lines<'a>(lines: &[&'a [u8]]) -> Vec<&'a [u8]> {
51 if let Some(last) = lines.iter().rposition(|line| !line.is_empty()) {
52 lines[0..=last].to_vec()
53 } else {
54 lines.to_vec()
55 }
56}
57
58pub fn cleanup_captured_buffer(buffer: &[u8], drop_n_last_lines: usize) -> Vec<u8> {
69 let trimmed_lines: Vec<&[u8]> = buf_trim_trailing(buffer);
70 let mut buffer: Vec<&[u8]> = drop_last_empty_lines(&trimmed_lines);
71 buffer.truncate(buffer.len() - drop_n_last_lines);
72
73 let mut final_buffer: Vec<u8> = Vec::with_capacity(buffer.len());
75 for (idx, &line) in buffer.iter().enumerate() {
76 final_buffer.extend_from_slice(line);
77
78 let is_last_line = idx == buffer.len() - 1;
79 if is_last_line {
80 let reset = "\u{001b}[0m".as_bytes();
81 final_buffer.extend_from_slice(reset);
82 final_buffer.push(b'\n');
83 } else {
84 final_buffer.push(b'\n');
85 }
86 }
87
88 final_buffer
89}
90
91#[cfg(test)]
92mod tests {
93 use super::{buf_trim_trailing, drop_last_empty_lines, SliceExt};
94
95 #[test]
96 fn trims_trailing_whitespaces() {
97 let input = " text ".as_bytes();
98 let expected = " text".as_bytes();
99
100 let actual = input.trim_trailing();
101 assert_eq!(actual, expected);
102 }
103
104 #[test]
105 fn trims_whitespaces() {
106 let input = " text ".as_bytes();
107 let expected = "text".as_bytes();
108
109 let actual = input.trim();
110 assert_eq!(actual, expected);
111 }
112
113 #[test]
114 fn test_buf_trim_trailing() {
115 let text = "line1\n\nline3 ";
116 let actual = buf_trim_trailing(text.as_bytes());
117 let expected = vec!["line1".as_bytes(), "".as_bytes(), "line3".as_bytes()];
118 assert_eq!(actual, expected);
119 }
120
121 #[test]
122 fn test_buf_drop_last_empty_lines() {
123 let text = "line1\nline2\n\nline3 ";
124
125 let trimmed_lines = buf_trim_trailing(text.as_bytes());
126 let actual = drop_last_empty_lines(&trimmed_lines);
127 let expected = trimmed_lines;
128 assert_eq!(actual, expected);
129
130 let text = "line1\nline2\n\n\n ";
133
134 let trimmed_lines = buf_trim_trailing(text.as_bytes());
135 let actual = drop_last_empty_lines(&trimmed_lines);
136 let expected = vec!["line1".as_bytes(), "line2".as_bytes()];
137 assert_eq!(actual, expected);
138 }
139}