1use crate::io::{self, BufWriter, IoSlice, Write};
2use crate::sys_common::memchr;
3
4/// Private helper struct for implementing the line-buffered writing logic.
5/// This shim temporarily wraps a BufWriter, and uses its internals to
6/// implement a line-buffered writer (specifically by using the internal
7/// methods like write_to_buf and flush_buf). In this way, a more
8/// efficient abstraction can be created than one that only had access to
9/// `write` and `flush`, without needlessly duplicating a lot of the
10/// implementation details of BufWriter. This also allows existing
11/// `BufWriters` to be temporarily given line-buffering logic; this is what
12/// enables Stdout to be alternately in line-buffered or block-buffered mode.
13#[derive(Debug)]
14pub struct LineWriterShim<'a, W: ?Sized + Write> {
15 buffer: &'a mut BufWriter<W>,
16}
17
18impl<'a, W: ?Sized + Write> LineWriterShim<'a, W> {
19 pub fn new(buffer: &'a mut BufWriter<W>) -> Self {
20 Self { buffer }
21 }
22
23 /// Get a reference to the inner writer (that is, the writer
24 /// wrapped by the BufWriter).
25 fn inner(&self) -> &W {
26 self.buffer.get_ref()
27 }
28
29 /// Get a mutable reference to the inner writer (that is, the writer
30 /// wrapped by the BufWriter). Be careful with this writer, as writes to
31 /// it will bypass the buffer.
32 fn inner_mut(&mut self) -> &mut W {
33 self.buffer.get_mut()
34 }
35
36 /// Get the content currently buffered in self.buffer
37 fn buffered(&self) -> &[u8] {
38 self.buffer.buffer()
39 }
40
41 /// Flush the buffer iff the last byte is a newline (indicating that an
42 /// earlier write only succeeded partially, and we want to retry flushing
43 /// the buffered line before continuing with a subsequent write)
44 fn flush_if_completed_line(&mut self) -> io::Result<()> {
45 match self.buffered().last().copied() {
46 Some(b'\n') => self.buffer.flush_buf(),
47 _ => Ok(()),
48 }
49 }
50}
51
52impl<'a, W: ?Sized + Write> Write for LineWriterShim<'a, W> {
53 /// Write some data into this BufReader with line buffering. This means
54 /// that, if any newlines are present in the data, the data up to the last
55 /// newline is sent directly to the underlying writer, and data after it
56 /// is buffered. Returns the number of bytes written.
57 ///
58 /// This function operates on a "best effort basis"; in keeping with the
59 /// convention of `Write::write`, it makes at most one attempt to write
60 /// new data to the underlying writer. If that write only reports a partial
61 /// success, the remaining data will be buffered.
62 ///
63 /// Because this function attempts to send completed lines to the underlying
64 /// writer, it will also flush the existing buffer if it ends with a
65 /// newline, even if the incoming data does not contain any newlines.
66 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
67 let newline_idx = match memchr::memrchr(b'\n', buf) {
68 // If there are no new newlines (that is, if this write is less than
69 // one line), just do a regular buffered write (which may flush if
70 // we exceed the inner buffer's size)
71 None => {
72 self.flush_if_completed_line()?;
73 return self.buffer.write(buf);
74 }
75 // Otherwise, arrange for the lines to be written directly to the
76 // inner writer.
77 Some(newline_idx) => newline_idx + 1,
78 };
79
80 // Flush existing content to prepare for our write. We have to do this
81 // before attempting to write `buf` in order to maintain consistency;
82 // if we add `buf` to the buffer then try to flush it all at once,
83 // we're obligated to return Ok(), which would mean suppressing any
84 // errors that occur during flush.
85 self.buffer.flush_buf()?;
86
87 // This is what we're going to try to write directly to the inner
88 // writer. The rest will be buffered, if nothing goes wrong.
89 let lines = &buf[..newline_idx];
90
91 // Write `lines` directly to the inner writer. In keeping with the
92 // `write` convention, make at most one attempt to add new (unbuffered)
93 // data. Because this write doesn't touch the BufWriter state directly,
94 // and the buffer is known to be empty, we don't need to worry about
95 // self.buffer.panicked here.
96 let flushed = self.inner_mut().write(lines)?;
97
98 // If buffer returns Ok(0), propagate that to the caller without
99 // doing additional buffering; otherwise we're just guaranteeing
100 // an "ErrorKind::WriteZero" later.
101 if flushed == 0 {
102 return Ok(0);
103 }
104
105 // Now that the write has succeeded, buffer the rest (or as much of
106 // the rest as possible). If there were any unwritten newlines, we
107 // only buffer out to the last unwritten newline that fits in the
108 // buffer; this helps prevent flushing partial lines on subsequent
109 // calls to LineWriterShim::write.
110
111 // Handle the cases in order of most-common to least-common, under
112 // the presumption that most writes succeed in totality, and that most
113 // writes are smaller than the buffer.
114 // - Is this a partial line (ie, no newlines left in the unwritten tail)
115 // - If not, does the data out to the last unwritten newline fit in
116 // the buffer?
117 // - If not, scan for the last newline that *does* fit in the buffer
118 let tail = if flushed >= newline_idx {
119 &buf[flushed..]
120 } else if newline_idx - flushed <= self.buffer.capacity() {
121 &buf[flushed..newline_idx]
122 } else {
123 let scan_area = &buf[flushed..];
124 let scan_area = &scan_area[..self.buffer.capacity()];
125 match memchr::memrchr(b'\n', scan_area) {
126 Some(newline_idx) => &scan_area[..newline_idx + 1],
127 None => scan_area,
128 }
129 };
130
131 let buffered = self.buffer.write_to_buf(tail);
132 Ok(flushed + buffered)
133 }
134
135 fn flush(&mut self) -> io::Result<()> {
136 self.buffer.flush()
137 }
138
139 /// Write some vectored data into this BufReader with line buffering. This
140 /// means that, if any newlines are present in the data, the data up to
141 /// and including the buffer containing the last newline is sent directly
142 /// to the inner writer, and the data after it is buffered. Returns the
143 /// number of bytes written.
144 ///
145 /// This function operates on a "best effort basis"; in keeping with the
146 /// convention of `Write::write`, it makes at most one attempt to write
147 /// new data to the underlying writer.
148 ///
149 /// Because this function attempts to send completed lines to the underlying
150 /// writer, it will also flush the existing buffer if it contains any
151 /// newlines.
152 ///
153 /// Because sorting through an array of `IoSlice` can be a bit convoluted,
154 /// This method differs from write in the following ways:
155 ///
156 /// - It attempts to write the full content of all the buffers up to and
157 /// including the one containing the last newline. This means that it
158 /// may attempt to write a partial line, that buffer has data past the
159 /// newline.
160 /// - If the write only reports partial success, it does not attempt to
161 /// find the precise location of the written bytes and buffer the rest.
162 ///
163 /// If the underlying vector doesn't support vectored writing, we instead
164 /// simply write the first non-empty buffer with `write`. This way, we
165 /// get the benefits of more granular partial-line handling without losing
166 /// anything in efficiency
167 fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
168 // If there's no specialized behavior for write_vectored, just use
169 // write. This has the benefit of more granular partial-line handling.
170 if !self.is_write_vectored() {
171 return match bufs.iter().find(|buf| !buf.is_empty()) {
172 Some(buf) => self.write(buf),
173 None => Ok(0),
174 };
175 }
176
177 // Find the buffer containing the last newline
178 let last_newline_buf_idx = bufs
179 .iter()
180 .enumerate()
181 .rev()
182 .find_map(|(i, buf)| memchr::memchr(b'\n', buf).map(|_| i));
183
184 // If there are no new newlines (that is, if this write is less than
185 // one line), just do a regular buffered write
186 let last_newline_buf_idx = match last_newline_buf_idx {
187 // No newlines; just do a normal buffered write
188 None => {
189 self.flush_if_completed_line()?;
190 return self.buffer.write_vectored(bufs);
191 }
192 Some(i) => i,
193 };
194
195 // Flush existing content to prepare for our write
196 self.buffer.flush_buf()?;
197
198 // This is what we're going to try to write directly to the inner
199 // writer. The rest will be buffered, if nothing goes wrong.
200 let (lines, tail) = bufs.split_at(last_newline_buf_idx + 1);
201
202 // Write `lines` directly to the inner writer. In keeping with the
203 // `write` convention, make at most one attempt to add new (unbuffered)
204 // data. Because this write doesn't touch the BufWriter state directly,
205 // and the buffer is known to be empty, we don't need to worry about
206 // self.panicked here.
207 let flushed = self.inner_mut().write_vectored(lines)?;
208
209 // If inner returns Ok(0), propagate that to the caller without
210 // doing additional buffering; otherwise we're just guaranteeing
211 // an "ErrorKind::WriteZero" later.
212 if flushed == 0 {
213 return Ok(0);
214 }
215
216 // Don't try to reconstruct the exact amount written; just bail
217 // in the event of a partial write
218 let lines_len = lines.iter().map(|buf| buf.len()).sum();
219 if flushed < lines_len {
220 return Ok(flushed);
221 }
222
223 // Now that the write has succeeded, buffer the rest (or as much of the
224 // rest as possible)
225 let buffered: usize = tail
226 .iter()
227 .filter(|buf| !buf.is_empty())
228 .map(|buf| self.buffer.write_to_buf(buf))
229 .take_while(|&n| n > 0)
230 .sum();
231
232 Ok(flushed + buffered)
233 }
234
235 fn is_write_vectored(&self) -> bool {
236 self.inner().is_write_vectored()
237 }
238
239 /// Write some data into this BufReader with line buffering. This means
240 /// that, if any newlines are present in the data, the data up to the last
241 /// newline is sent directly to the underlying writer, and data after it
242 /// is buffered.
243 ///
244 /// Because this function attempts to send completed lines to the underlying
245 /// writer, it will also flush the existing buffer if it contains any
246 /// newlines, even if the incoming data does not contain any newlines.
247 fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
248 match memchr::memrchr(b'\n', buf) {
249 // If there are no new newlines (that is, if this write is less than
250 // one line), just do a regular buffered write (which may flush if
251 // we exceed the inner buffer's size)
252 None => {
253 self.flush_if_completed_line()?;
254 self.buffer.write_all(buf)
255 }
256 Some(newline_idx) => {
257 let (lines, tail) = buf.split_at(newline_idx + 1);
258
259 if self.buffered().is_empty() {
260 self.inner_mut().write_all(lines)?;
261 } else {
262 // If there is any buffered data, we add the incoming lines
263 // to that buffer before flushing, which saves us at least
264 // one write call. We can't really do this with `write`,
265 // since we can't do this *and* not suppress errors *and*
266 // report a consistent state to the caller in a return
267 // value, but here in write_all it's fine.
268 self.buffer.write_all(lines)?;
269 self.buffer.flush_buf()?;
270 }
271
272 self.buffer.write_all(tail)
273 }
274 }
275 }
276}
277