use crate::io::{self, IoSlice, Write};
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct VectoredWriter<'io, const N: usize> {
skip: usize,
data: [&'io [u8]; N],
}
#[allow(dead_code)]
impl<'io, const N: usize> VectoredWriter<'io, N> {
pub const fn new(data: [&'io [u8]; N]) -> Self {
Self { skip: 0, data }
}
pub fn write_vectored_inner(
&mut self,
bufs: &mut [IoSlice<'io>; N],
writer: &mut impl Write,
) -> io::Result<usize> {
if self.skip >= N {
return Ok(0);
}
let len = writer.write_vectored(&bufs[self.skip..])?;
let mut remove = 0;
let mut accumulated_len = 0;
for buf in &self.data[self.skip..] {
if accumulated_len + buf.len() > len {
break;
}
accumulated_len += buf.len();
remove += 1;
}
self.skip += remove;
if self.skip < N {
let rem = len - accumulated_len;
self.data[self.skip] = &self.data[self.skip][rem..];
bufs[self.skip] = IoSlice::new(self.data[self.skip]);
}
Ok(len)
}
pub fn write_all_vectored(&mut self, mut writer: impl Write) -> io::Result<usize> {
if self.skip >= N {
return Ok(0);
}
let mut bufs: [IoSlice; N] = self.data.map(IoSlice::new);
let mut written = 0;
while {
let len = self.write_vectored_inner(&mut bufs, &mut writer)?;
if len == 0 {
return Ok(written);
}
written += len;
self.skip < N
} {}
Ok(written)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[test]
fn test_with_split_writer() {
struct SplitWriter {
data: Cursor<Vec<u8>>,
cnt: usize,
reset: usize,
}
impl Write for SplitWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
let wrt = if self.cnt != 0 && self.cnt < buf.len() {
println!("Partial write: {:?}", &buf[0..self.cnt]);
self.data.write(&buf[0..self.cnt])?
} else {
println!("Complete write: {buf:?}");
self.data.write(buf)?
};
self.cnt = self.cnt.saturating_sub(wrt);
if self.cnt == 0 {
self.cnt = self.reset;
}
Ok(wrt)
}
fn flush(&mut self) -> std::io::Result<()> {
self.data.flush()
}
}
let data: [&[u8]; 4] = [b"dead", b"beef", b"cafe", b"babe"];
let mut out = SplitWriter {
data: Cursor::new(Vec::with_capacity(16)),
cnt: 5,
reset: 5,
};
println!("Preparing to write {data:?}");
let mut vecw = VectoredWriter::new(data);
println!("{vecw:?}");
let n = vecw.write_all_vectored(&mut out).unwrap();
assert_eq!(n, 16);
assert_eq!(&out.data.into_inner()[..], b"deadbeefcafebabe");
}
}