use crate::config::{FileServeConfig, FsTaskSpawner};
use crate::file::ChunkedReadFile;
use std::io;
use core::{mem, task, future};
use core::pin::Pin;
pub enum Body<W: FsTaskSpawner, C: FileServeConfig> {
Full(bytes::Bytes),
Chunked(ChunkedReadFile<W, C>)
}
impl<W: FsTaskSpawner, C: FileServeConfig> Body<W, C> {
#[inline(always)]
pub const fn empty() -> Self {
Self::Full(bytes::Bytes::new())
}
pub fn poll_next(self: Pin<&mut Self>, ctx: &mut task::Context<'_>) -> task::Poll<Option<Result<bytes::Bytes, io::Error>>> {
match self.get_mut() {
Self::Full(bytes) => {
if bytes.is_empty() {
task::Poll::Ready(None)
} else {
let mut result = bytes::Bytes::new();
mem::swap(&mut result, bytes);
task::Poll::Ready(Some(Ok(result)))
}
},
Self::Chunked(chunked) => match future::Future::poll(Pin::new(chunked), ctx) {
task::Poll::Pending => task::Poll::Pending,
task::Poll::Ready(Ok(Some(result))) => {
task::Poll::Ready(Some(Ok(result)))
}
task::Poll::Ready(Ok(None)) => task::Poll::Ready(None),
task::Poll::Ready(Err(error)) => task::Poll::Ready(Some(Err(error))),
}
}
}
#[inline(always)]
pub fn is_finished(&self) -> bool {
match self {
Self::Full(bytes) => bytes.is_empty(),
Self::Chunked(chunked) => chunked.is_finished(),
}
}
#[inline(always)]
pub fn len(&self) -> u64 {
match self {
Self::Full(bytes) => bytes.len() as u64,
Self::Chunked(chunked) => chunked.remaining()
}
}
}
impl<W: FsTaskSpawner, C: FileServeConfig> From<bytes::Bytes> for Body<W, C> {
#[inline(always)]
fn from(bytes: bytes::Bytes) -> Self {
Body::Full(bytes)
}
}
impl<W: FsTaskSpawner, C: FileServeConfig> From<ChunkedReadFile<W, C>> for Body<W, C> {
#[inline(always)]
fn from(chunked: ChunkedReadFile<W, C>) -> Self {
Body::Chunked(chunked)
}
}
impl<W: FsTaskSpawner, C: FileServeConfig> From<Option<ChunkedReadFile<W, C>>> for Body<W, C> {
#[inline(always)]
fn from(chunked: Option<ChunkedReadFile<W, C>>) -> Self {
match chunked {
Some(chunked) => chunked.into(),
None => Self::empty()
}
}
}