use crate::{Error, Result};
pub trait EncodeWriter {
fn write_all(&mut self, bytes: &[u8]) -> Result<()>;
}
impl<W: EncodeWriter + ?Sized> EncodeWriter for &mut W {
fn write_all(&mut self, bytes: &[u8]) -> Result<()> {
(**self).write_all(bytes)
}
}
#[derive(Debug)]
pub struct SliceWriter<'a> {
output: &'a mut [u8],
required: usize,
}
impl<'a> SliceWriter<'a> {
pub fn new(output: &'a mut [u8]) -> Self {
Self {
output,
required: 0,
}
}
pub const fn required_len(&self) -> usize {
self.required
}
pub fn written_len(&self) -> usize {
self.required.min(self.output.len())
}
pub fn finish(self) -> Result<usize> {
if self.required > self.output.len() {
Err(Error::BufferTooSmall {
required: self.required,
available: self.output.len(),
})
} else {
Ok(self.required)
}
}
}
impl EncodeWriter for SliceWriter<'_> {
fn write_all(&mut self, bytes: &[u8]) -> Result<()> {
let start = self.required.min(self.output.len());
let writable = (self.output.len() - start).min(bytes.len());
self.output[start..start + writable].copy_from_slice(&bytes[..writable]);
self.required = self
.required
.checked_add(bytes.len())
.ok_or(Error::SizeLimit { limit: u64::MAX })?;
Ok(())
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct CountWriter {
written: u64,
}
impl CountWriter {
pub const fn new() -> Self {
Self { written: 0 }
}
pub const fn written(&self) -> u64 {
self.written
}
}
impl EncodeWriter for CountWriter {
fn write_all(&mut self, bytes: &[u8]) -> Result<()> {
let amount =
u64::try_from(bytes.len()).map_err(|_| Error::IntegerOverflow { target: "u64" })?;
self.written = self
.written
.checked_add(amount)
.ok_or(Error::SizeLimit { limit: u64::MAX })?;
Ok(())
}
}
#[cfg(feature = "alloc")]
impl EncodeWriter for alloc::vec::Vec<u8> {
fn write_all(&mut self, bytes: &[u8]) -> Result<()> {
self.extend_from_slice(bytes);
Ok(())
}
}