use crate::error::Result;
use bytes::Bytes;
use futures::Stream;
use std::fmt;
use std::pin::Pin;
pub type ByteStream<'a> = Pin<Box<dyn Stream<Item = Result<Bytes>> + Send + 'a>>;
#[derive(Clone, Copy)]
pub struct MmapView<'a>(&'a [u8]);
impl<'a> MmapView<'a> {
#[must_use]
pub fn new(bytes: &'a [u8]) -> Self {
Self(bytes)
}
#[must_use]
pub fn as_slice(&self) -> &'a [u8] {
self.0
}
}
impl fmt::Debug for MmapView<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MmapView")
.field("len", &self.0.len())
.finish()
}
}
impl AsRef<[u8]> for MmapView<'_> {
fn as_ref(&self) -> &[u8] {
self.0
}
}
pub enum DocPayload<'a> {
Borrowed(&'a [u8]),
Mapped(MmapView<'a>),
Owned(Bytes),
Stream(ByteStream<'a>),
}
impl DocPayload<'_> {
#[must_use]
pub fn as_slice(&self) -> Option<&[u8]> {
match self {
Self::Borrowed(b) => Some(b),
Self::Mapped(m) => Some(m.as_slice()),
Self::Owned(b) => Some(b),
Self::Stream(_) => None,
}
}
#[must_use]
pub fn len(&self) -> Option<usize> {
self.as_slice().map(<[u8]>::len)
}
#[must_use]
pub fn is_empty(&self) -> Option<bool> {
self.as_slice().map(<[u8]>::is_empty)
}
pub async fn into_bytes(self) -> Result<Bytes> {
match self {
Self::Borrowed(b) => Ok(Bytes::copy_from_slice(b)),
Self::Mapped(m) => Ok(Bytes::copy_from_slice(m.as_slice())),
Self::Owned(b) => Ok(b),
Self::Stream(mut s) => {
use futures::StreamExt;
let mut buf = Vec::new();
while let Some(chunk) = s.next().await {
buf.extend_from_slice(&chunk?);
}
Ok(Bytes::from(buf))
}
}
}
}
impl fmt::Debug for DocPayload<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Borrowed(b) => f.debug_tuple("Borrowed").field(&b.len()).finish(),
Self::Mapped(m) => f.debug_tuple("Mapped").field(&m.as_slice().len()).finish(),
Self::Owned(b) => f.debug_tuple("Owned").field(&b.len()).finish(),
Self::Stream(_) => f.write_str("Stream(..)"),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn borrowed_and_owned_expose_slices() {
let data = b"hello world";
let borrowed = DocPayload::Borrowed(data);
assert_eq!(borrowed.as_slice(), Some(&data[..]));
assert_eq!(borrowed.len(), Some(11));
assert_eq!(borrowed.is_empty(), Some(false));
let owned = DocPayload::Owned(Bytes::from_static(data));
assert_eq!(owned.as_slice(), Some(&data[..]));
}
#[test]
fn mapped_view_roundtrips() {
let data = [1u8, 2, 3, 4];
let view = MmapView::new(&data);
assert_eq!(view.as_slice(), &data);
assert_eq!(view.as_ref(), &data);
}
#[test]
fn stream_has_no_in_memory_slice() {
let s: ByteStream<'_> = Box::pin(futures::stream::empty());
let payload = DocPayload::Stream(s);
assert_eq!(payload.as_slice(), None);
assert_eq!(payload.len(), None);
}
#[tokio::test]
async fn into_bytes_drains_a_stream() {
let chunks = vec![
Ok(Bytes::from_static(b"foo")),
Ok(Bytes::from_static(b"bar")),
];
let s: ByteStream<'_> = Box::pin(futures::stream::iter(chunks));
let bytes = DocPayload::Stream(s).into_bytes().await.unwrap();
assert_eq!(&bytes[..], b"foobar");
}
#[tokio::test]
async fn into_bytes_on_borrowed_copies() {
let bytes = DocPayload::Borrowed(b"abc").into_bytes().await.unwrap();
assert_eq!(&bytes[..], b"abc");
}
}