use crate::types::ErrorCode;
use crate::types::Result;
pub const BUFFER_MAX_SIZE: usize = 64 * 1024 * 1024;
pub const BUFFER_INITIAL_SIZE: usize = 4096;
pub const BUFFER_RESET_CAPACITY: usize = BUFFER_INITIAL_SIZE;
const BUFFER_GROW_FACTOR: usize = 2;
#[derive(Debug)]
pub struct Buffer {
data: Vec<u8>,
size: usize,
read_pos: usize,
owned: bool,
}
impl Default for Buffer {
fn default() -> Self {
Self::new()
}
}
impl Buffer {
pub fn new() -> Self {
Self {
data: vec![0u8; BUFFER_INITIAL_SIZE],
size: 0,
read_pos: 0,
owned: true,
}
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
data: vec![0u8; capacity],
size: 0,
read_pos: 0,
owned: true,
}
}
pub fn from_slice(data: &[u8]) -> Self {
Self {
data: data.to_vec(),
size: data.len(),
read_pos: 0,
owned: true,
}
}
pub fn from_static(data: &mut [u8]) -> Buffer {
let len = data.len();
Buffer {
data: data.to_vec(),
size: len,
read_pos: 0,
owned: false,
}
}
pub fn available(&self) -> usize {
self.size.saturating_sub(self.read_pos)
}
pub fn space(&self) -> usize {
self.data.len().saturating_sub(self.size)
}
pub fn write(&mut self, data: &[u8]) -> Result<usize> {
if data.is_empty() {
return Ok(0);
}
self.compact();
if data.len() > BUFFER_MAX_SIZE || self.size > BUFFER_MAX_SIZE - data.len() {
return Err(ErrorCode::Internal);
}
let needed = self.size + data.len();
self.ensure_capacity(needed)?;
self.data[self.size..self.size + data.len()].copy_from_slice(data);
self.size += data.len();
Ok(data.len())
}
pub fn read(&mut self, out: &mut [u8]) -> Result<usize> {
if self.read_pos + out.len() > self.size {
return Err(ErrorCode::Io);
}
out.copy_from_slice(&self.data[self.read_pos..self.read_pos + out.len()]);
self.read_pos += out.len();
Ok(out.len())
}
pub fn peek(&self) -> &[u8] {
&self.data[self.read_pos..self.size]
}
pub fn capacity(&self) -> usize {
self.data.len()
}
pub fn reset(&mut self) {
self.size = 0;
self.read_pos = 0;
if self.owned && self.data.len() > BUFFER_RESET_CAPACITY {
self.data.truncate(BUFFER_RESET_CAPACITY);
self.data.shrink_to_fit();
}
}
pub fn drain(&mut self, len: usize) {
let available = self.available();
self.read_pos += len.min(available);
}
pub fn as_slice(&self) -> &[u8] {
&self.data[..self.size]
}
pub fn as_mut_slice(&mut self) -> &mut [u8] {
&mut self.data[..self.size]
}
pub fn write_pos(&self) -> usize {
self.size
}
pub fn set_size(&mut self, size: usize) {
self.size = size.min(self.data.len());
}
pub fn read_pos(&self) -> usize {
self.read_pos
}
pub fn set_read_pos(&mut self, pos: usize) {
self.read_pos = pos.min(self.size);
}
fn compact(&mut self) {
if self.read_pos == 0 {
return;
}
if self.read_pos >= self.size {
self.size = 0;
self.read_pos = 0;
return;
}
let avail = self.size - self.read_pos;
self.data.copy_within(self.read_pos..self.size, 0);
self.size = avail;
self.read_pos = 0;
}
fn ensure_capacity(&mut self, needed: usize) -> Result<()> {
if needed <= self.data.len() {
return Ok(());
}
if !self.owned {
return Err(ErrorCode::Internal);
}
if needed > BUFFER_MAX_SIZE {
return Err(ErrorCode::Internal);
}
let mut new_cap = self.data.len().max(BUFFER_INITIAL_SIZE);
while new_cap < needed {
new_cap = new_cap.saturating_mul(BUFFER_GROW_FACTOR);
if new_cap > BUFFER_MAX_SIZE {
new_cap = BUFFER_MAX_SIZE;
break;
}
}
self.data.resize(new_cap, 0);
Ok(())
}
}
impl Clone for Buffer {
fn clone(&self) -> Self {
Self {
data: self.data.clone(),
size: self.size,
read_pos: self.read_pos,
owned: true,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn write_then_read_round_trips() {
let mut buf = Buffer::new();
buf.write(b"hello world").unwrap();
assert_eq!(buf.available(), 11);
let mut out = [0u8; 5];
buf.read(&mut out).unwrap();
assert_eq!(&out, b"hello");
assert_eq!(buf.available(), 6);
assert_eq!(buf.peek(), b" world");
}
#[test]
fn read_past_available_errors() {
let mut buf = Buffer::new();
buf.write(b"ab").unwrap();
let mut out = [0u8; 3];
assert!(buf.read(&mut out).is_err());
}
#[test]
fn drain_skips_unread_bytes() {
let mut buf = Buffer::new();
buf.write(b"abcdef").unwrap();
buf.drain(3);
assert_eq!(buf.peek(), b"def");
}
#[test]
fn grows_past_initial_capacity() {
let mut buf = Buffer::new();
let chunk = vec![0xAB_u8; BUFFER_INITIAL_SIZE * 3];
buf.write(&chunk).unwrap();
assert_eq!(buf.available(), chunk.len());
assert_eq!(buf.peek(), chunk.as_slice());
}
#[test]
fn reset_clears_buffer() {
let mut buf = Buffer::new();
buf.write(b"data").unwrap();
buf.reset();
assert_eq!(buf.available(), 0);
assert_eq!(buf.peek(), b"");
}
#[test]
fn reset_shrinks_oversized_capacity() {
let mut buf = Buffer::new();
buf.write(&vec![0u8; 1024 * 1024]).unwrap();
assert!(buf.capacity() >= 1024 * 1024);
buf.reset();
assert_eq!(buf.available(), 0);
assert!(buf.capacity() <= BUFFER_RESET_CAPACITY);
}
#[test]
fn write_over_max_size_errors() {
let mut buf = Buffer::new();
assert!(buf.write(&vec![0u8; BUFFER_MAX_SIZE + 1]).is_err());
}
}