toasty_core/stmt/

value_stream.rs

use crate::schema::Schema;
use crate::stmt::Type;

use super::Value;

use std::{
    collections::VecDeque,
    fmt, mem,
    panic::Location,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};
use tokio_stream::{Stream, StreamExt};

/// An async stream of [`Value`]s with optional type checking.
///
/// `ValueStream` combines a buffered front-end with an optional async
/// [`Stream`] back-end. Values can be pushed into the buffer or pulled
/// from the underlying stream. When a [`Type`] is attached via
/// [`typed`](ValueStream::typed), every yielded value is checked at
/// runtime.
///
/// Implements [`Stream`] from `tokio_stream`, yielding
/// `Result<Value>` items.
///
/// # Examples
///
/// ```ignore
/// use toasty_core::stmt::{Value, ValueStream};
///
/// let mut stream = ValueStream::from_value(Value::from(42_i64));
/// let val = stream.next().await.unwrap().unwrap();
/// assert_eq!(val, Value::from(42_i64));
/// ```
#[derive(Default)]
pub struct ValueStream {
    buffer: Buffer,
    stream: Option<DynStream>,

    /// If set, check values to ensure they are the correct type. The schema
    /// resolves `Type::Model` (`#[document]`) field layouts for the check.
    ty: Option<(Arc<Schema>, Type, &'static Location<'static>)>,
}

#[derive(Debug)]
struct Iter<I> {
    iter: I,
}

#[derive(Clone, Default)]
enum Buffer {
    #[default]
    Empty,
    One(Value),
    Many(VecDeque<Value>),
}

type DynStream = Pin<Box<dyn Stream<Item = crate::Result<Value>> + Send + 'static>>;

impl ValueStream {
    /// Creates a stream containing a single value.
    pub fn from_value(value: impl Into<Value>) -> Self {
        Self {
            buffer: Buffer::One(value.into()),
            stream: None,
            ty: None,
        }
    }

    /// Creates a stream backed by an async [`Stream`] of `Result<Value>`.
    pub fn from_stream<T: Stream<Item = crate::Result<Value>> + Send + 'static>(stream: T) -> Self {
        Self {
            buffer: Buffer::Empty,
            stream: Some(Box::pin(stream)),
            ty: None,
        }
    }

    /// Creates a fully-buffered stream from a vector of values.
    pub fn from_vec(records: Vec<Value>) -> Self {
        Self {
            buffer: Buffer::Many(records.into()),
            stream: None,
            ty: None,
        }
    }

    /// Creates a stream from a fallible iterator.
    #[allow(clippy::should_implement_trait)]
    pub fn from_iter<T, I>(iter: I) -> Self
    where
        T: Into<Value>,
        I: Iterator<Item = crate::Result<T>> + Send + 'static,
    {
        Self::from_stream(Iter { iter })
    }

    /// Returns the next record in the stream
    pub async fn next(&mut self) -> Option<crate::Result<Value>> {
        StreamExt::next(self).await
    }

    /// Peek at the next record in the stream
    pub async fn peek(&mut self) -> Option<crate::Result<&Value>> {
        if self.buffer.is_empty() {
            match self.next().await {
                Some(Ok(value)) => self.buffer.push(value),
                Some(Err(e)) => return Some(Err(e)),
                None => return None,
            }
        }

        self.buffer.first().map(Ok)
    }

    /// Force the stream to preload at least one record, if there are more
    /// records to stream.
    pub async fn tap(&mut self) -> crate::Result<()> {
        if let Some(Err(e)) = self.peek().await {
            Err(e)
        } else {
            Ok(())
        }
    }

    /// Returns the minimum number of elements this stream will yield.
    ///
    /// This is derived from the stream's `size_hint` lower bound plus
    /// the number of buffered elements.
    pub fn min_len(&self) -> usize {
        let (ret, _) = self.size_hint();
        ret
    }

    /// Consumes the stream and collects all values into a `Vec`.
    pub async fn collect(mut self) -> crate::Result<Vec<Value>> {
        let mut ret = Vec::with_capacity(self.min_len());

        while let Some(res) = self.next().await {
            ret.push(res?);
        }

        Ok(ret)
    }

    /// Fully buffers the stream and returns a clone of it.
    ///
    /// After this call, both the original and the returned stream are
    /// fully buffered and contain the same values.
    pub async fn dup(&mut self) -> crate::Result<Self> {
        self.buffer().await?;

        Ok(Self {
            buffer: self.buffer.clone(),
            stream: None,
            ty: self.ty.clone(),
        })
    }

    /// Returns a clone if the stream is fully buffered, or `None` if
    /// there is an unconsumed async stream that cannot be cloned.
    pub fn try_clone(&self) -> Option<Self> {
        if self.stream.is_some() {
            return None;
        }

        Some(Self {
            buffer: self.buffer.clone(),
            stream: None,
            ty: self.ty.clone(),
        })
    }

    /// Drains the underlying async stream into the buffer.
    ///
    /// After this call, all remaining values are buffered locally and
    /// [`is_buffered`](ValueStream::is_buffered) returns `true`.
    pub async fn buffer(&mut self) -> crate::Result<()> {
        if let Some(stream) = &mut self.stream {
            while let Some(res) = stream.next().await {
                let value = res?;

                if let Some((schema, ty, location)) = &self.ty {
                    assert!(
                        value.is_a(&schema.app, ty),
                        "expected `{ty:?}`; was={value:#?}; origin={location}"
                    );
                }

                self.buffer.push(value);
            }
        }

        Ok(())
    }

    /// Returns `true` if the ValueStream is fully buffered (no remaining stream)
    pub fn is_buffered(&self) -> bool {
        self.stream.is_none()
    }

    /// Returns a clone of only the currently buffered values
    /// Does not consume any stream data or wait for additional values
    pub fn buffered_to_vec(&self) -> Vec<Value> {
        match &self.buffer {
            Buffer::Empty => Vec::new(),
            Buffer::One(value) => vec![value.clone()],
            Buffer::Many(values) => values.iter().cloned().collect(),
        }
    }

    /// Returns a mutable iterator over the buffered values.
    ///
    /// # Panics
    ///
    /// Panics if the stream has an unconsumed async back-end. Call
    /// [`buffer`](ValueStream::buffer) first to ensure all values are
    /// buffered.
    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut Value> {
        assert!(self.stream.is_none());

        // TODO: don't box
        match &mut self.buffer {
            Buffer::Empty => Box::new(None.into_iter()),
            Buffer::One(v) => Box::new(Some(v).into_iter()),
            Buffer::Many(v) => Box::new(v.iter_mut()) as Box<dyn Iterator<Item = &mut Value>>,
        }
    }

    /// Attaches a [`Type`] constraint to this stream.
    ///
    /// Every value yielded from the stream (both already-buffered and
    /// future) will be checked against `ty` at runtime. If a value does
    /// not match, the check panics with a diagnostic message.
    ///
    /// # Panics
    ///
    /// Panics if an already-buffered value is not compatible with `ty`,
    /// or if a previously set type differs from `ty`.
    #[track_caller]
    pub fn typed(mut self, schema: Arc<Schema>, ty: Type) -> ValueStream {
        let location = Location::caller();

        match &self.ty {
            Some((_, prev, _)) => assert_eq!(*prev, ty),
            None => {
                // Validate already-buffered values against the new type.
                // Document values have already been raised to their positional
                // `Value::Record` form by the engine (see the engine's
                // document raising), so this is a shape check only.
                let mut tmp = mem::take(&mut self.buffer);
                while let Some(value) = tmp.next() {
                    assert!(
                        value.is_a(&schema.app, &ty),
                        "expected `{ty:?}`; was={value:#?}; origin={location}"
                    );
                    self.buffer.push(value);
                }

                self.ty = Some((schema, ty, location));
            }
        }

        self
    }
}

impl Stream for ValueStream {
    type Item = crate::Result<Value>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if let Some(next) = self.buffer.next() {
            Poll::Ready(Some(Ok(next)))
        } else if let Some(stream) = self.stream.as_mut() {
            match Pin::new(stream).poll_next(cx) {
                Poll::Ready(Some(Ok(value))) => {
                    if let Some((schema, ty, location)) = &self.ty {
                        assert!(
                            value.is_a(&schema.app, ty),
                            "expected `{ty:?}`; was={value:#?}; origin={location}"
                        );
                    }
                    Poll::Ready(Some(Ok(value)))
                }

                other => other,
            }
        } else {
            Poll::Ready(None)
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (mut low, mut high) = match &self.stream {
            Some(stream) => stream.size_hint(),
            None => (0, Some(0)),
        };

        let buffered = self.buffer.len();

        low += buffered;

        if let Some(high) = high.as_mut() {
            *high += buffered;
        }

        (low, high)
    }
}

impl From<Value> for ValueStream {
    fn from(src: Value) -> Self {
        Self {
            buffer: Buffer::One(src),
            stream: None,
            ty: None,
        }
    }
}

impl From<Vec<Value>> for ValueStream {
    fn from(value: Vec<Value>) -> Self {
        Self::from_vec(value)
    }
}

impl<I> Unpin for Iter<I> {}

impl<T, I> Stream for Iter<I>
where
    I: Iterator<Item = crate::Result<T>>,
    T: Into<Value>,
{
    type Item = crate::Result<Value>;

    fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Poll::Ready(self.iter.next().map(|res| res.map(|item| item.into())))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl fmt::Debug for ValueStream {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RecordStream").finish()
    }
}

impl Buffer {
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn len(&self) -> usize {
        match self {
            Self::Empty => 0,
            Self::One(_) => 1,
            Self::Many(v) => v.len(),
        }
    }

    fn first(&self) -> Option<&Value> {
        match self {
            Self::Empty => None,
            Self::One(value) => Some(value),
            Self::Many(values) => values.front(),
        }
    }

    fn next(&mut self) -> Option<Value> {
        match self {
            Self::Empty => None,
            Self::One(_) => {
                let Self::One(value) = mem::take(self) else {
                    panic!()
                };
                Some(value)
            }
            Self::Many(values) => values.pop_front(),
        }
    }

    fn push(&mut self, value: Value) {
        match self {
            Self::Empty => {
                *self = Self::One(value);
            }
            Self::One(_) => {
                let Self::One(first) = mem::replace(self, Self::Many(VecDeque::with_capacity(2)))
                else {
                    panic!()
                };

                let Self::Many(values) = self else { panic!() };

                values.push_back(first);
                values.push_back(value);
            }
            Self::Many(values) => {
                values.push_back(value);
            }
        }
    }
}