use std::collections::VecDeque;
use std::pin::Pin;
use std::task::{Context, Poll};
use futures_core::Stream;
use super::ByteStream;
use crate::error::{Error, Result};
#[derive(Debug)]
pub struct LineStream {
inner: ByteStream,
buf: Vec<u8>,
ready: VecDeque<String>,
done: bool,
}
impl LineStream {
pub(super) fn new(inner: ByteStream) -> Self {
Self {
inner,
buf: Vec::new(),
ready: VecDeque::new(),
done: false,
}
}
fn drain_lines(&mut self) -> Result<()> {
while let Some(nl) = self.buf.iter().position(|&b| b == b'\n') {
let mut line: Vec<u8> = self.buf.drain(..=nl).collect();
line.pop(); if line.last() == Some(&b'\r') {
line.pop();
}
let line =
String::from_utf8(line).map_err(|e| Error::decode("invalid utf-8 line", e))?;
self.ready.push_back(line);
}
Ok(())
}
}
impl Stream for LineStream {
type Item = Result<String>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
loop {
if let Some(line) = self.ready.pop_front() {
return Poll::Ready(Some(Ok(line)));
}
if self.done {
if self.buf.is_empty() {
return Poll::Ready(None);
}
let rest = std::mem::take(&mut self.buf);
return match String::from_utf8(rest) {
Ok(s) => Poll::Ready(Some(Ok(s))),
Err(e) => Poll::Ready(Some(Err(Error::decode("invalid utf-8 line", e)))),
};
}
match Pin::new(&mut self.inner).poll_next(cx) {
Poll::Ready(Some(Ok(chunk))) => {
self.buf.extend_from_slice(&chunk);
if let Err(e) = self.drain_lines() {
return Poll::Ready(Some(Err(e)));
}
}
Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))),
Poll::Ready(None) => self.done = true,
Poll::Pending => return Poll::Pending,
}
}
}
}