use crate::error::{EncodeError, EncodeResult};
pub trait Encodable {
fn encode(&self, e: &mut Encoder<'_>) -> EncodeResult<()>;
fn to_vec(&self) -> EncodeResult<Vec<u8>> {
let mut buffer = Vec::with_capacity(512);
let mut encoder = Encoder::new(&mut buffer);
self.encode(&mut encoder)?;
Ok(buffer)
}
}
#[derive(Debug)]
pub struct Encoder<'a> {
buffer: &'a mut Vec<u8>,
offset: usize,
}
impl<'a> Encoder<'a> {
pub fn new(buffer: &'a mut Vec<u8>) -> Self {
Self { buffer, offset: 0 }
}
pub fn buffer(&self) -> &[u8] {
self.buffer
}
pub fn write_slice(&mut self, bytes: &[u8]) -> EncodeResult<()> {
let additional = bytes.len();
if self.offset + additional <= self.buffer.len() {
for (byte, b) in self.buffer[self.offset..].iter_mut().zip(bytes.iter()) {
*byte = *b;
}
} else {
let expected_len = self.buffer.len() + additional;
self.buffer.reserve(additional);
self.buffer.extend_from_slice(bytes);
debug_assert!(self.buffer.len() == expected_len);
}
let index = self
.offset
.checked_add(additional)
.ok_or(EncodeError::AddOverflow)?;
self.offset = index;
Ok(())
}
pub fn write<const N: usize>(&mut self, bytes: [u8; N]) -> EncodeResult<()> {
let additional = bytes.len();
if self.offset + additional <= self.buffer.len() {
for (byte, b) in self.buffer[self.offset..].iter_mut().zip(bytes) {
*byte = b;
}
} else {
let expected_len = self.buffer.len() + additional;
self.buffer.reserve(additional);
self.buffer.extend(bytes);
debug_assert!(self.buffer.len() == expected_len);
}
let index = self
.offset
.checked_add(additional)
.ok_or(EncodeError::AddOverflow)?;
self.offset = index;
Ok(())
}
pub fn write_u8(&mut self, data: u8) -> EncodeResult<()> {
self.write(data.to_be_bytes())
}
pub fn write_u16(&mut self, data: u16) -> EncodeResult<()> {
self.write(data.to_be_bytes())
}
pub fn write_u32(&mut self, data: u32) -> EncodeResult<()> {
self.write(data.to_be_bytes())
}
pub fn write_u128(&mut self, data: u128) -> EncodeResult<()> {
self.write(data.to_be_bytes())
}
pub fn write_u64(&mut self, data: u64) -> EncodeResult<()> {
self.write(data.to_be_bytes())
}
pub fn write_i32(&mut self, data: i32) -> EncodeResult<()> {
self.write(data.to_be_bytes())
}
pub fn write_fill_bytes(&mut self, bytes: &[u8], fill_len: usize) -> EncodeResult<()> {
if bytes.len() > fill_len {
return Err(EncodeError::StringSizeTooBig { len: bytes.len() });
}
let nul_len = fill_len - bytes.len();
self.write_slice(bytes)?;
for _ in 0..nul_len {
self.write_u8(0)?;
}
Ok(())
}
pub fn write_fill<T: AsRef<[u8]>>(
&mut self,
s: &Option<T>,
fill_len: usize,
) -> EncodeResult<()> {
match s {
Some(sname) => {
let bytes = sname.as_ref();
self.write_fill_bytes(bytes, fill_len)?;
}
None => {
for _ in 0..fill_len {
self.write_u8(0)?;
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn basic_encode() -> EncodeResult<()> {
let mut buf = vec![0, 1, 2, 3, 4, 5];
let mut enc = Encoder::new(&mut buf);
enc.offset = 4;
enc.write_slice(&[5, 6])?;
assert_eq!(enc.buffer, &mut vec![0, 1, 2, 3, 5, 6]);
assert_eq!(enc.offset, 6);
enc.write_slice(&[7, 8])?;
assert_eq!(enc.buffer, &mut vec![0, 1, 2, 3, 5, 6, 7, 8]);
assert_eq!(enc.offset, 8);
let mut buf = vec![];
let mut enc = Encoder::new(&mut buf);
enc.write_slice(&[0, 1, 2, 3])?;
assert_eq!(enc.buffer, &mut vec![0, 1, 2, 3]);
assert_eq!(enc.offset, 4);
Ok(())
}
}