francoisgib_webserver 1.0.3

mod test;

use std::{io::BufRead, str::FromStr};
use tokio::io::{AsyncBufRead, AsyncReadExt};

use crate::config::PRE_ALLOCATED_BUFFER_SIZE;

/// Represents the two possible storage strategies for a `Buffer`.
#[derive(Debug, Clone)]
pub enum BufferType {
    /// Small data fits on the stack
    Stack([u8; PRE_ALLOCATED_BUFFER_SIZE]),
    /// Larger data is stored on the heap
    Heap(Vec<u8>),
}

/// Represents a buffer with known read size.
///
/// - Optimized to use stack if small, and heap for larger content.
/// - Provides sync and async reads.
/// - Implements `FromStr` and `ToString`.
#[derive(Debug, Clone)]
pub struct Buffer {
    buffer: BufferType,
    read_size: usize,
}

impl Buffer {
    /// Creates a new buffer with the desired size.
    pub fn new(size: usize) -> Self {
        Self {
            buffer: BufferType::new(size),
            read_size: 0,
        }
    }

    /// Synchronously reads into the buffer from a `BufRead` reader.
    pub fn read<R: BufRead>(&mut self, reader: &mut R) -> Result<usize, String> {
        match self.buffer.read(reader) {
            Ok(bytes_read) => {
                self.read_size += bytes_read;
                Ok(bytes_read)
            }
            Err(e) => Err(e),
        }
    }

    /// Asynchronously reads into the buffer from an `AsyncBufRead` reader.
    pub async fn async_read<R: AsyncBufRead + Unpin>(
        &mut self,
        reader: &mut R,
    ) -> Result<usize, String> {
        match self.buffer.async_read(reader).await {
            Ok(bytes_read) => {
                self.read_size += bytes_read;
                Ok(bytes_read)
            }
            Err(e) => Err(e),
        }
    }

    /// Returns the number of bytes currently stored.
    pub fn len(&self) -> usize {
        self.read_size
    }

    pub fn to_bytes(self) -> Vec<u8> {
        match self.buffer {
            BufferType::Heap(h) => h,
            BufferType::Stack(s) => s[..self.read_size].to_vec(),
        }
    }
}

impl ToString for Buffer {
    fn to_string(&self) -> String {
        match &self.buffer {
            BufferType::Stack(buffer) => {
                String::from_utf8_lossy(&buffer[..self.read_size]).to_string()
            }
            BufferType::Heap(vec) => String::from_utf8_lossy(&vec[..self.read_size]).to_string(),
        }
    }
}

impl FromStr for Buffer {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(Self {
            buffer: BufferType::from_str(s)?,
            read_size: s.len(),
        })
    }
}

impl BufferType {
    /// Chooses `Stack` or `Heap` based on the requested size.
    pub fn new(size: usize) -> Self {
        if size <= PRE_ALLOCATED_BUFFER_SIZE && size > 0 {
            Self::Stack([0u8; PRE_ALLOCATED_BUFFER_SIZE])
        } else {
            let mut vec = Vec::with_capacity(size);
            vec.resize(vec.capacity(), 0);
            Self::Heap(vec)
        }
    }

    /// Synchronously reads into the buffer.
    pub fn read<R: BufRead>(&mut self, reader: &mut R) -> Result<usize, String> {
        match self {
            Self::Stack(buffer) => reader.read(buffer).map_err(|e| e.to_string()),
            Self::Heap(vec) => reader.read(vec).map_err(|e| e.to_string()),
        }
    }

    /// Asynchronously reads into the buffer.
    pub async fn async_read<R: AsyncBufRead + Unpin>(
        &mut self,
        reader: &mut R,
    ) -> Result<usize, String> {
        match self {
            Self::Stack(buffer) => reader.read(buffer).await.map_err(|e| e.to_string()),
            Self::Heap(vec) => reader.read(vec).await.map_err(|e| e.to_string()),
        }
    }
}

impl FromStr for BufferType {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut buffer = BufferType::new(s.len());
        let bytes = s.as_bytes();
        match &mut buffer {
            Self::Stack(arr) => arr[..bytes.len()].copy_from_slice(bytes),
            Self::Heap(vec) => vec.extend_from_slice(bytes),
        }
        Ok(buffer)
    }
}

impl From<Vec<u8>> for Buffer {
    fn from(vec: Vec<u8>) -> Self {
        let read_size = vec.len();
        let buffer = if read_size <= PRE_ALLOCATED_BUFFER_SIZE {
            let mut stack_buffer = [0u8; PRE_ALLOCATED_BUFFER_SIZE];
            stack_buffer[..read_size].copy_from_slice(&vec);
            BufferType::Stack(stack_buffer)
        } else {
            BufferType::Heap(vec)
        };

        Self { buffer, read_size }
    }
}