1 //! Insert line breaks between written buffers when they would overflow the line length.
2 use 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.
7 pub(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 
15 impl<W: io::Write> AutoBreak<W> {
new(writer: W, line_capacity: usize) -> Self16     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 
flush_buf(&mut self) -> io::Result<()>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 
58 impl<W: io::Write> io::Write for AutoBreak<W> {
write(&mut self, buffer: &[u8]) -> io::Result<usize>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 
flush(&mut self) -> io::Result<()>76     fn flush(&mut self) -> io::Result<()> {
77         self.flush_buf()?;
78         self.wrapped.flush()
79     }
80 }
81 
82 impl<W: io::Write> Drop for AutoBreak<W> {
drop(&mut self)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)]
92 mod tests {
93     use super::*;
94     use std::io::Write;
95 
96     #[test]
test_aligned_writes()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]
test_greater_writes()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 }
125