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