use crate::error::DecodeError;
use crate::leb128;
pub struct FrameWriter {
buf: Vec<u8>,
}
impl FrameWriter {
pub fn new() -> Self {
Self {
buf: Vec::with_capacity(256),
}
}
pub fn write_frame(&mut self, payload: &[u8]) {
leb128::encode(&mut self.buf, payload.len() as u64);
self.buf.extend_from_slice(payload);
}
pub fn as_bytes(&self) -> &[u8] {
&self.buf
}
pub fn finish(self) -> Vec<u8> {
self.buf
}
pub fn reset(&mut self) {
self.buf.clear();
}
}
impl Default for FrameWriter {
fn default() -> Self {
Self::new()
}
}
pub struct FrameReader<'a> {
data: &'a [u8],
pos: usize,
}
impl<'a> FrameReader<'a> {
pub fn new(data: &'a [u8]) -> Self {
Self { data, pos: 0 }
}
pub fn read_frame(&mut self) -> Option<Result<&'a [u8], DecodeError>> {
if self.pos >= self.data.len() {
return None;
}
let (len, consumed) = match leb128::decode(&self.data[self.pos..], 10) {
Ok(v) => v,
Err(e) => return Some(Err(e)),
};
self.pos += consumed;
let len = len as usize;
if self.pos + len > self.data.len() {
return Some(Err(DecodeError::UnexpectedEof));
}
let payload = &self.data[self.pos..self.pos + len];
self.pos += len;
Some(Ok(payload))
}
pub fn is_empty(&self) -> bool {
self.pos >= self.data.len()
}
pub fn remaining(&self) -> &'a [u8] {
&self.data[self.pos..]
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{BitReader, BitWriter};
#[test]
fn roundtrip_single_frame() {
let mut w = BitWriter::new();
w.write_u32(42);
let payload = w.finish();
let mut fw = FrameWriter::new();
fw.write_frame(&payload);
let stream = fw.finish();
let mut fr = FrameReader::new(&stream);
let frame = fr.read_frame().unwrap().unwrap();
let mut r = BitReader::new(frame);
assert_eq!(r.read_u32().unwrap(), 42);
assert!(fr.read_frame().is_none());
}
#[test]
fn roundtrip_multiple_frames() {
let mut fw = FrameWriter::new();
let mut w = BitWriter::new();
w.write_u32(1);
fw.write_frame(&w.finish());
let mut w = BitWriter::new();
w.write_string("hello");
fw.write_frame(&w.finish());
let mut w = BitWriter::new();
w.write_u64(999);
fw.write_frame(&w.finish());
let stream = fw.finish();
let mut fr = FrameReader::new(&stream);
let f1 = fr.read_frame().unwrap().unwrap();
assert_eq!(BitReader::new(f1).read_u32().unwrap(), 1);
let f2 = fr.read_frame().unwrap().unwrap();
assert_eq!(BitReader::new(f2).read_string().unwrap(), "hello");
let f3 = fr.read_frame().unwrap().unwrap();
assert_eq!(BitReader::new(f3).read_u64().unwrap(), 999);
assert!(fr.read_frame().is_none());
}
#[test]
fn writer_reset_reuse() {
let mut fw = FrameWriter::new();
let mut w = BitWriter::new();
w.write_u32(1);
fw.write_frame(&w.finish());
let bytes1 = fw.finish();
let mut fw = FrameWriter::new();
let mut w = BitWriter::new();
w.write_u32(1);
fw.write_frame(&w.finish());
let bytes2 = fw.finish();
assert_eq!(bytes1, bytes2);
}
#[test]
fn empty_stream() {
let mut fr = FrameReader::new(&[]);
assert!(fr.read_frame().is_none());
assert!(fr.is_empty());
}
#[test]
fn corrupted_length() {
let stream = [0x80]; let mut fr = FrameReader::new(&stream);
match fr.read_frame() {
Some(Err(_)) => {} other => panic!("expected error, got {:?}", other),
}
}
#[test]
fn truncated_payload() {
let stream = [10, 0x01, 0x02, 0x03];
let mut fr = FrameReader::new(&stream);
match fr.read_frame() {
Some(Err(DecodeError::UnexpectedEof)) => {} other => panic!("expected EOF error, got {:?}", other),
}
}
}