use super::{BodyFrame, ReadableBody};
use base_scheduler::{Cancellation, Cancelled};
#[derive(Clone, Debug)]
pub struct BufferedBody<D = Vec<u8>> {
chunks: Chunks<D>,
index: usize,
length: Option<u64>,
}
#[derive(Clone, Debug)]
enum Chunks<D> {
Empty,
One(D),
Many(Vec<D>),
}
impl<D> Chunks<D> {
fn as_slice(&self) -> &[D] {
match self {
Self::Empty => &[],
Self::One(chunk) => std::slice::from_ref(chunk),
Self::Many(chunks) => chunks,
}
}
fn push(&mut self, chunk: D) {
match self {
Self::Empty => *self = Self::One(chunk),
Self::Many(chunks) => chunks.push(chunk),
Self::One(_) => {
let Self::One(first) = std::mem::replace(self, Self::Empty) else {
unreachable!();
};
*self = Self::Many(vec![first, chunk]);
}
}
}
}
impl<D: AsRef<[u8]> + Clone + Send> BufferedBody<D> {
pub fn new(chunks: impl IntoIterator<Item = D>) -> Self {
let mut iter = chunks.into_iter();
let chunks = match iter.next() {
None => Chunks::Empty,
Some(first) => match iter.next() {
None => Chunks::One(first),
Some(second) => {
let mut chunks = vec![first, second];
chunks.extend(iter);
Chunks::Many(chunks)
}
},
};
let length = chunks.as_slice().iter().try_fold(0u64, |sum, chunk| {
u64::try_from(chunk.as_ref().len())
.ok()
.and_then(|length| sum.checked_add(length))
});
Self {
chunks,
index: 0,
length,
}
}
pub fn chunks(&self) -> &[D] {
self.chunks.as_slice()
}
pub fn read_chunk(&mut self) -> Option<D> {
let chunk = self.chunks.as_slice().get(self.index)?.clone();
self.index += 1;
Some(chunk)
}
pub fn buffered(&self) -> Self {
Self {
chunks: self.chunks.clone(),
index: 0,
length: self.length,
}
}
pub fn write(&mut self, chunk: D) {
self.length = self
.length
.zip(u64::try_from(chunk.as_ref().len()).ok())
.and_then(|(length, added)| length.checked_add(added));
self.chunks.push(chunk);
}
pub fn clear(&mut self) {
self.chunks = Chunks::Empty;
self.index = 0;
self.length = Some(0);
}
}
impl<D> Default for BufferedBody<D> {
fn default() -> Self {
Self {
chunks: Chunks::Empty,
index: 0,
length: Some(0),
}
}
}
impl<D: AsRef<[u8]> + Clone + Send> From<D> for BufferedBody<D> {
fn from(chunk: D) -> Self {
Self::new([chunk])
}
}
impl<D: AsRef<[u8]> + Clone + Send> ReadableBody for BufferedBody<D> {
type Data = D;
type Error = Cancelled;
fn may_have_trailers(&self) -> bool {
false
}
fn is_ready(&self) -> bool {
true
}
fn is_empty(&self) -> bool {
self.index >= self.chunks.as_slice().len()
}
fn length(&self) -> Option<u64> {
self.length
}
fn is_rewindable(&self) -> bool {
true
}
fn rewind(&mut self) -> bool {
self.index = 0;
true
}
fn close(&mut self, _error: Option<Self::Error>) {
self.index = self.chunks.as_slice().len();
}
async fn read(
&mut self,
cancellation: &Cancellation,
) -> Result<BodyFrame<Self::Data>, Self::Error> {
cancellation.check()?;
Ok(match self.read_chunk() {
Some(data) => BodyFrame::Data {
data,
end: self.is_empty(),
},
None => BodyFrame::End,
})
}
}