1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

/// Misc utilities.

pub(crate) trait SliceExt {
    fn trim(&self) -> &Self;
    fn trim_trailing(&self) -> &Self;
}

fn is_whitespace(c: &u8) -> bool {
    *c == b'\t' || *c == b' '
}

fn is_not_whitespace(c: &u8) -> bool {
    !is_whitespace(c)
}

impl SliceExt for [u8] {
    /// Trim leading and trailing whitespaces (`\t` and ` `) in a `&[u8]`
    fn trim(&self) -> &[u8] {
        if let Some(first) = self.iter().position(is_not_whitespace) {
            if let Some(last) = self.iter().rposition(is_not_whitespace) {
                &self[first..last + 1]
            } else {
                unreachable!();
            }
        } else {
            &[]
        }
    }

    /// Trim trailing whitespaces (`\t` and ` `) in a `&[u8]`
    fn trim_trailing(&self) -> &[u8] {
        if let Some(last) = self.iter().rposition(is_not_whitespace) {
            &self[0..last + 1]
        } else {
            &[]
        }
    }
}

/// Trim each line of the buffer.
fn buf_trim_trailing(buf: &[u8]) -> Vec<&[u8]> {
    let trimmed_lines: Vec<&[u8]> = buf
        .split(|c| *c == b'\n')
        .map(SliceExt::trim_trailing) // trim each line
        .collect();

    trimmed_lines
}

/// Drop all the last empty lines.
fn drop_last_empty_lines<'a>(lines: &[&'a [u8]]) -> Vec<&'a [u8]> {
    if let Some(last) = lines.iter().rposition(|line| !line.is_empty()) {
        lines[0..=last].to_vec()
    } else {
        lines.to_vec()
    }
}

/// This function processes a pane captured bufer.
///
/// - All lines are trimmed after capture because tmux does not allow capturing escape codes and
///   trimming lines.
/// - If `drop_n_last_lines` is greater than 0, the n last line are not captured. This is used only
///   for panes with a zsh prompt, in order to avoid polluting the history with new prompts on
///   restore.
/// - In addition, the last line has an additional ascii reset escape code because tmux does not
///   capture it.
///
pub fn cleanup_captured_buffer(buffer: &[u8], drop_n_last_lines: usize) -> Vec<u8> {
    let trimmed_lines: Vec<&[u8]> = buf_trim_trailing(buffer);
    let mut buffer: Vec<&[u8]> = drop_last_empty_lines(&trimmed_lines);
    buffer.truncate(buffer.len() - drop_n_last_lines);

    // Join the lines with `b'\n'`, add reset code to the last line
    let mut final_buffer: Vec<u8> = Vec::with_capacity(buffer.len());
    for (idx, &line) in buffer.iter().enumerate() {
        final_buffer.extend_from_slice(line);

        let is_last_line = idx == buffer.len() - 1;
        if is_last_line {
            let reset = "\u{001b}[0m".as_bytes();
            final_buffer.extend_from_slice(reset);
            final_buffer.push(b'\n');
        } else {
            final_buffer.push(b'\n');
        }
    }

    final_buffer
}

#[cfg(test)]
mod tests {
    use super::{buf_trim_trailing, drop_last_empty_lines, SliceExt};

    #[test]
    fn trims_trailing_whitespaces() {
        let input = "  text   ".as_bytes();
        let expected = "  text".as_bytes();

        let actual = input.trim_trailing();
        assert_eq!(actual, expected);
    }

    #[test]
    fn trims_whitespaces() {
        let input = "  text   ".as_bytes();
        let expected = "text".as_bytes();

        let actual = input.trim();
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_buf_trim_trailing() {
        let text = "line1\n\nline3   ";
        let actual = buf_trim_trailing(text.as_bytes());
        let expected = vec!["line1".as_bytes(), "".as_bytes(), "line3".as_bytes()];
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_buf_drop_last_empty_lines() {
        let text = "line1\nline2\n\nline3   ";

        let trimmed_lines = buf_trim_trailing(text.as_bytes());
        let actual = drop_last_empty_lines(&trimmed_lines);
        let expected = trimmed_lines;
        assert_eq!(actual, expected);

        //

        let text = "line1\nline2\n\n\n     ";

        let trimmed_lines = buf_trim_trailing(text.as_bytes());
        let actual = drop_last_empty_lines(&trimmed_lines);
        let expected = vec!["line1".as_bytes(), "line2".as_bytes()];
        assert_eq!(actual, expected);
    }
}