Simplify send.rs

1 files changed, 18 insertions, 11 deletions

diff --git a/src/stream.rs b/src/stream.rs
index 6c23f4a..ae57d62 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -82,6 +82,23 @@ impl ByteStreamReader {
 		}
 	}
 
+	pub fn take_buffer(&mut self) -> Bytes {
+		let bytes = Bytes::from(
+			self .buf
+				.iter()
+				.map(|x| &x[..])
+				.collect::<Vec<_>>()
+				.concat(),
+		);
+		self.buf.clear();
+		self.buf_len = 0;
+		bytes
+	}
+
+	pub fn eos(&self) -> bool {
+		self.buf.is_empty() && self.eos
+	}
+
 	fn try_get(&mut self, read_len: usize) -> Option<Bytes> {
 		if self.buf_len >= read_len {
 			let mut slices = Vec::with_capacity(self.buf.len());
@@ -144,17 +161,7 @@ impl<'a> Future for ByteStreamReadExact<'a> {
 				if *this.fail_on_eos {
 					return Poll::Ready(Err(ReadExactError::UnexpectedEos));
 				} else {
-					let bytes = Bytes::from(
-						this.reader
-							.buf
-							.iter()
-							.map(|x| &x[..])
-							.collect::<Vec<_>>()
-							.concat(),
-					);
-					this.reader.buf.clear();
-					this.reader.buf_len = 0;
-					return Poll::Ready(Ok(bytes));
+					return Poll::Ready(Ok(this.reader.take_buffer()));
 				}
 			}