image/codecs/pnm/
autobreak.rs

1//! Insert line breaks between written buffers when they would overflow the line length.
2use std::io;
3
4// The pnm standard says to insert line breaks after 70 characters. Assumes that no line breaks
5// are actually written. We have to be careful to fully commit buffers or not commit them at all,
6// otherwise we might insert a newline in the middle of a token.
7pub(crate) struct AutoBreak<W: io::Write> {
8    wrapped: W,
9    line_capacity: usize,
10    line: Vec<u8>,
11    has_newline: bool,
12    panicked: bool, // see https://github.com/rust-lang/rust/issues/30888
13}
14
15impl<W: io::Write> AutoBreak<W> {
16    pub(crate) fn new(writer: W, line_capacity: usize) -> Self {
17        AutoBreak {
18            wrapped: writer,
19            line_capacity,
20            line: Vec::with_capacity(line_capacity + 1),
21            has_newline: false,
22            panicked: false,
23        }
24    }
25
26    fn flush_buf(&mut self) -> io::Result<()> {
27        // from BufWriter
28        let mut written = 0;
29        let len = self.line.len();
30        let mut ret = Ok(());
31        while written < len {
32            self.panicked = true;
33            let r = self.wrapped.write(&self.line[written..]);
34            self.panicked = false;
35            match r {
36                Ok(0) => {
37                    ret = Err(io::Error::new(
38                        io::ErrorKind::WriteZero,
39                        "failed to write the buffered data",
40                    ));
41                    break;
42                }
43                Ok(n) => written += n,
44                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
45                Err(e) => {
46                    ret = Err(e);
47                    break;
48                }
49            }
50        }
51        if written > 0 {
52            self.line.drain(..written);
53        }
54        ret
55    }
56}
57
58impl<W: io::Write> io::Write for AutoBreak<W> {
59    fn write(&mut self, buffer: &[u8]) -> io::Result<usize> {
60        if self.has_newline {
61            self.flush()?;
62            self.has_newline = false;
63        }
64
65        if !self.line.is_empty() && self.line.len() + buffer.len() > self.line_capacity {
66            self.line.push(b'\n');
67            self.has_newline = true;
68            self.flush()?;
69            self.has_newline = false;
70        }
71
72        self.line.extend_from_slice(buffer);
73        Ok(buffer.len())
74    }
75
76    fn flush(&mut self) -> io::Result<()> {
77        self.flush_buf()?;
78        self.wrapped.flush()
79    }
80}
81
82impl<W: io::Write> Drop for AutoBreak<W> {
83    fn drop(&mut self) {
84        if !self.panicked {
85            let _r = self.flush_buf();
86            // internal writer flushed automatically by Drop
87        }
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use std::io::Write;
95
96    #[test]
97    fn test_aligned_writes() {
98        let mut output = Vec::new();
99
100        {
101            let mut writer = AutoBreak::new(&mut output, 10);
102            writer.write_all(b"0123456789").unwrap();
103            writer.write_all(b"0123456789").unwrap();
104        }
105
106        assert_eq!(output.as_slice(), b"0123456789\n0123456789");
107    }
108
109    #[test]
110    fn test_greater_writes() {
111        let mut output = Vec::new();
112
113        {
114            let mut writer = AutoBreak::new(&mut output, 10);
115            writer.write_all(b"012").unwrap();
116            writer.write_all(b"345").unwrap();
117            writer.write_all(b"0123456789").unwrap();
118            writer.write_all(b"012345678910").unwrap();
119            writer.write_all(b"_").unwrap();
120        }
121
122        assert_eq!(output.as_slice(), b"012345\n0123456789\n012345678910\n_");
123    }
124}