diff --git a/crates/async-compression/src/generic/write/buf_writer.rs b/crates/async-compression/src/generic/write/buf_writer.rs index 3dc05591..22f677e5 100644 --- a/crates/async-compression/src/generic/write/buf_writer.rs +++ b/crates/async-compression/src/generic/write/buf_writer.rs @@ -47,7 +47,7 @@ impl BufWriter { } /// Remove the already written data - fn reshuffle_and_remove_written(&mut self) { + fn remove_written(&mut self) { self.buf.copy_within(self.written..self.buffered, 0); self.buffered -= self.written; self.written = 0; @@ -82,11 +82,7 @@ impl BufWriter { Ok(()) }; - if self.written > 0 { - self.reshuffle_and_remove_written(); - - Poll::Ready(ret) - } else if self.buffered == 0 { + if self.written > 0 || self.buffered < self.buf.len() { Poll::Ready(ret) } else { ret?; @@ -99,7 +95,11 @@ impl BufWriter { poll_write: &mut dyn FnMut(&[u8]) -> Poll>, ) -> Poll> { let ret = ready!(self.do_flush(poll_write)); - self.reshuffle_and_remove_written(); + + debug_assert_eq!(self.buffered, self.written); + self.buffered = 0; + self.written = 0; + Poll::Ready(ret) } @@ -108,22 +108,24 @@ impl BufWriter { buf: &[u8], poll_write: &mut dyn FnMut(&[u8]) -> Poll>, ) -> Poll> { - if self.buffered + buf.len() > self.buf.len() { - ready!(self.partial_flush_buf(poll_write))?; - } - if buf.len() >= self.buf.len() { - if self.buffered == 0 { - poll_write(buf) - } else { - // The only way that `partial_flush_buf` would have returned with - // `this.buffered != 0` is if it were Pending, so our waker was already queued - Poll::Pending - } + ready!(self.flush_buf(poll_write))?; + poll_write(buf) + } else if (self.buf.len() - self.buffered) >= buf.len() { + self.buf[self.buffered..].copy_from_slice(buf); + self.buffered += buf.len(); + + Poll::Ready(Ok(buf.len())) } else { + ready!(self.partial_flush_buf(poll_write))?; + if self.written > 0 { + self.remove_written(); + } + let len = buf.len().min(self.buf.len() - self.buffered); self.buf[self.buffered..self.buffered + len].copy_from_slice(&buf[..len]); self.buffered += len; + Poll::Ready(Ok(len)) } } @@ -133,6 +135,11 @@ impl BufWriter { poll_write: &mut dyn FnMut(&[u8]) -> Poll>, ) -> Poll> { ready!(self.partial_flush_buf(poll_write))?; + + if self.written > 0 { + self.remove_written(); + } + Poll::Ready(Ok(&mut self.buf[self.buffered..])) }