1 files changed, 60 insertions, 56 deletions

diff --git a/src/send.rs b/src/send.rs
index 330d41d..0179eb2 100644
--- a/src/send.rs
+++ b/src/send.rs
@@ -24,6 +24,7 @@ use crate::util::{ByteStream, Packet};
 
 pub(crate) type RequestID = u32;
 pub(crate) type ChunkLength = u16;
+
 pub(crate) const MAX_CHUNK_LENGTH: ChunkLength = 0x3FF0;
 pub(crate) const ERROR_MARKER: ChunkLength = 0x4000;
 pub(crate) const CHUNK_HAS_CONTINUATION: ChunkLength = 0x8000;
@@ -62,53 +63,58 @@ enum DataFrame {
 		data: [u8; MAX_CHUNK_LENGTH as usize],
 		/// actual lenght of data
 		len: usize,
+		/// whethere there may be more data comming from this stream. Can be used for some
+		/// optimization. It's an error to set it to false if there is more data, but it is correct
+		/// (albeit sub-optimal) to set it to true if there is nothing coming after
+		may_have_more: bool,
 	},
+	/// An error code automatically signals the end of the stream
 	Error(u8),
 }
 
-struct DataReaderItem {
-	data: DataFrame,
-	/// whethere there may be more data comming from this stream. Can be used for some
-	/// optimization. It's an error to set it to false if there is more data, but it is correct
-	/// (albeit sub-optimal) to set it to true if there is nothing coming after
-	may_have_more: bool,
-}
-
-impl DataReaderItem {
+impl DataFrame {
 	fn empty_last() -> Self {
-		DataReaderItem {
-			data: DataFrame::Data {
-				data: [0; MAX_CHUNK_LENGTH as usize],
-				len: 0,
-			},
+		DataFrame::Data {
+			data: [0; MAX_CHUNK_LENGTH as usize],
+			len: 0,
 			may_have_more: false,
 		}
 	}
 
 	fn header(&self) -> [u8; 2] {
-		let continuation = if self.may_have_more {
-			CHUNK_HAS_CONTINUATION
-		} else {
-			0
-		};
-		let len = match self.data {
-			DataFrame::Data { len, .. } => len as u16,
-			DataFrame::Error(e) => e as u16 | ERROR_MARKER,
+		let header_u16 = match self {
+			DataFrame::Data {
+				len,
+				may_have_more: false,
+				..
+			} => *len as u16,
+			DataFrame::Data {
+				len,
+				may_have_more: true,
+				..
+			} => *len as u16 | CHUNK_HAS_CONTINUATION,
+			DataFrame::Error(e) => *e as u16 | ERROR_MARKER,
 		};
-
-		ChunkLength::to_be_bytes(len | continuation)
+		ChunkLength::to_be_bytes(header_u16)
 	}
 
 	fn data(&self) -> &[u8] {
-		match self.data {
-			DataFrame::Data { ref data, len } => &data[..len],
+		match self {
+			DataFrame::Data { ref data, len, .. } => &data[..*len],
 			DataFrame::Error(_) => &[],
 		}
 	}
+
+	fn may_have_more(&self) -> bool {
+		match self {
+			DataFrame::Data { may_have_more, .. } => *may_have_more,
+			DataFrame::Error(_) => false,
+		}
+	}
 }
 
 impl Stream for DataReader {
-	type Item = DataReaderItem;
+	type Item = DataFrame;
 
 	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
 		let mut this = self.project();
@@ -125,10 +131,8 @@ impl Stream for DataReader {
 				Err(e) => {
 					let e = *e;
 					*this.packet = Ok(Vec::new());
-					return Poll::Ready(Some(DataReaderItem {
-						data: DataFrame::Error(e),
-						may_have_more: true,
-					}));
+					*this.eos = true;
+					return Poll::Ready(Some(DataFrame::Error(e)));
 				}
 			};
 			let packet_left = packet.len() - *this.pos;
@@ -161,8 +165,9 @@ impl Stream for DataReader {
 		let len = this.buf.len();
 		body[..len].copy_from_slice(this.buf);
 		this.buf.clear();
-		Poll::Ready(Some(DataReaderItem {
-			data: DataFrame::Data { data: body, len },
+		Poll::Ready(Some(DataFrame::Data {
+			data: body,
+			len,
 			may_have_more: !*this.eos,
 		}))
 	}
@@ -218,38 +223,37 @@ struct SendQueuePollNextReady<'a> {
 }
 
 impl<'a> futures::Future for SendQueuePollNextReady<'a> {
-	type Output = (RequestID, DataReaderItem);
+	type Output = (RequestID, DataFrame);
 
 	fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
 		for i in 0..self.queue.items.len() {
 			let (_prio, items_at_prio) = &mut self.queue.items[i];
 
-			for _ in 0..items_at_prio.len() {
-				let mut item = items_at_prio.pop_front().unwrap();
-
+			let mut ready_item = None;
+			for (j, item) in items_at_prio.iter_mut().enumerate() {
 				match Pin::new(&mut item.data).poll_next(ctx) {
-					Poll::Pending => items_at_prio.push_back(item),
-					Poll::Ready(Some(data)) => {
-						let id = item.id;
-						if data.may_have_more {
-							self.queue.push(item);
-						} else {
-							if items_at_prio.is_empty() {
-								// this priority level is empty, remove it
-								self.queue.items.remove(i);
-							}
-						}
-						return Poll::Ready((id, data));
-					}
-					Poll::Ready(None) => {
-						if items_at_prio.is_empty() {
-							// this priority level is empty, remove it
-							self.queue.items.remove(i);
-						}
-						return Poll::Ready((item.id, DataReaderItem::empty_last()));
+					Poll::Pending => (),
+					Poll::Ready(ready_v) => {
+						ready_item = Some((j, ready_v));
+						break;
 					}
 				}
 			}
+
+			if let Some((j, ready_v)) = ready_item {
+				let item = items_at_prio.remove(j).unwrap();
+				let id = item.id;
+				if ready_v
+					.as_ref()
+					.map(|data| data.may_have_more())
+					.unwrap_or(false)
+				{
+					items_at_prio.push_back(item);
+				} else if items_at_prio.is_empty() {
+					self.queue.items.remove(i);
+				}
+				return Poll::Ready((id, ready_v.unwrap_or_else(DataFrame::empty_last)));
+			}
 		}
 		// TODO what do we do if self.queue is empty? We won't get scheduled again.
 		Poll::Pending