use alloc::collections::VecDeque;
use alloc::vec::Vec;
use core::cmp;
use std::io;
use std::io::Read;
pub(crate) struct ChunkVecBuffer {
chunks: VecDeque<Vec<u8>>,
limit: Option<usize>,
}
impl ChunkVecBuffer {
pub(crate) fn new(limit: Option<usize>) -> Self {
Self {
chunks: VecDeque::new(),
limit,
}
}
pub(crate) fn set_limit(&mut self, new_limit: Option<usize>) {
self.limit = new_limit;
}
pub(crate) fn is_empty(&self) -> bool {
self.chunks.is_empty()
}
pub(crate) fn is_full(&self) -> bool {
self.limit
.map(|limit| self.len() > limit)
.unwrap_or_default()
}
pub(crate) fn len(&self) -> usize {
let mut len = 0;
for ch in &self.chunks {
len += ch.len();
}
len
}
pub(crate) fn apply_limit(&self, len: usize) -> usize {
if let Some(limit) = self.limit {
let space = limit.saturating_sub(self.len());
cmp::min(len, space)
} else {
len
}
}
pub(crate) fn append_limited_copy(&mut self, bytes: &[u8]) -> usize {
let take = self.apply_limit(bytes.len());
self.append(bytes[..take].to_vec());
take
}
pub(crate) fn append(&mut self, bytes: Vec<u8>) -> usize {
let len = bytes.len();
if !bytes.is_empty() {
self.chunks.push_back(bytes);
}
len
}
pub(crate) fn pop(&mut self) -> Option<Vec<u8>> {
self.chunks.pop_front()
}
pub(crate) fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let mut offs = 0;
while offs < buf.len() && !self.is_empty() {
let used = self.chunks[0]
.as_slice()
.read(&mut buf[offs..])?;
self.consume(used);
offs += used;
}
Ok(offs)
}
#[cfg(read_buf)]
pub(crate) fn read_buf(&mut self, mut cursor: core::io::BorrowedCursor<'_>) -> io::Result<()> {
while !self.is_empty() && cursor.capacity() > 0 {
let chunk = self.chunks[0].as_slice();
let used = cmp::min(chunk.len(), cursor.capacity());
cursor.append(&chunk[..used]);
self.consume(used);
}
Ok(())
}
fn consume(&mut self, mut used: usize) {
while let Some(mut buf) = self.chunks.pop_front() {
if used < buf.len() {
buf.drain(..used);
self.chunks.push_front(buf);
break;
} else {
used -= buf.len();
}
}
}
pub(crate) fn write_to(&mut self, wr: &mut dyn io::Write) -> io::Result<usize> {
if self.is_empty() {
return Ok(0);
}
let mut bufs = [io::IoSlice::new(&[]); 64];
for (iov, chunk) in bufs.iter_mut().zip(self.chunks.iter()) {
*iov = io::IoSlice::new(chunk);
}
let len = cmp::min(bufs.len(), self.chunks.len());
let used = wr.write_vectored(&bufs[..len])?;
self.consume(used);
Ok(used)
}
}
#[cfg(test)]
mod tests {
use super::ChunkVecBuffer;
#[test]
fn short_append_copy_with_limit() {
let mut cvb = ChunkVecBuffer::new(Some(12));
assert_eq!(cvb.append_limited_copy(b"hello"), 5);
assert_eq!(cvb.append_limited_copy(b"world"), 5);
assert_eq!(cvb.append_limited_copy(b"hello"), 2);
assert_eq!(cvb.append_limited_copy(b"world"), 0);
let mut buf = [0u8; 12];
assert_eq!(cvb.read(&mut buf).unwrap(), 12);
assert_eq!(buf.to_vec(), b"helloworldhe".to_vec());
}
#[cfg(read_buf)]
#[test]
fn read_buf() {
use core::io::BorrowedBuf;
use core::mem::MaybeUninit;
{
let mut cvb = ChunkVecBuffer::new(None);
cvb.append(b"test ".to_vec());
cvb.append(b"fixture ".to_vec());
cvb.append(b"data".to_vec());
let mut buf = [MaybeUninit::<u8>::uninit(); 8];
let mut buf: BorrowedBuf<'_> = buf.as_mut_slice().into();
cvb.read_buf(buf.unfilled()).unwrap();
assert_eq!(buf.filled(), b"test fix");
buf.clear();
cvb.read_buf(buf.unfilled()).unwrap();
assert_eq!(buf.filled(), b"ture dat");
buf.clear();
cvb.read_buf(buf.unfilled()).unwrap();
assert_eq!(buf.filled(), b"a");
}
{
let mut cvb = ChunkVecBuffer::new(None);
cvb.append(b"short message".to_vec());
let mut buf = [MaybeUninit::<u8>::uninit(); 1024];
let mut buf: BorrowedBuf<'_> = buf.as_mut_slice().into();
cvb.read_buf(buf.unfilled()).unwrap();
assert_eq!(buf.filled(), b"short message");
}
}
}