rust-cache-server 0.1.0

use std::fmt;

use bytes::{Buf, BytesMut};
use thiserror::Error;

#[derive(Debug, Error)]
pub enum RespError {
    #[error("incomplete frame")]
    Incomplete,
    #[error("protocol error: {0}")]
    Protocol(String),
}

#[derive(Debug, Clone)]
pub enum RespFrame {
    SimpleString(String),
    Error(String),
    Integer(i64),
    BulkString(Option<Vec<u8>>),
    Array(Option<Vec<RespFrame>>),
}

impl fmt::Display for RespFrame {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            RespFrame::SimpleString(s) => write!(f, "{}", s),
            RespFrame::Error(s) => write!(f, "ERR {}", s),
            RespFrame::Integer(i) => write!(f, "{}", i),
            RespFrame::BulkString(Some(b)) => write!(f, "{}", String::from_utf8_lossy(b)),
            RespFrame::BulkString(None) => write!(f, "(nil)"),
            RespFrame::Array(Some(arr)) => {
                write!(f, "[")?;
                for (i, a) in arr.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{a}")?;
                }
                write!(f, "]")
            }
            RespFrame::Array(None) => write!(f, "(nil array)"),
        }
    }
}

fn read_line(buf: &mut BytesMut) -> Result<Option<Vec<u8>>, RespError> {
    if let Some(pos) = buf.windows(2).position(|w| w == b"\r\n") {
        let mut line = buf.split_to(pos + 2);
        line.truncate(pos);
        Ok(Some(line.to_vec()))
    } else {
        Ok(None)
    }
}

fn parse_isize(line: &[u8]) -> Result<isize, RespError> {
    let s = std::str::from_utf8(line).map_err(|e| RespError::Protocol(e.to_string()))?;
    s.parse::<isize>()
        .map_err(|e| RespError::Protocol(e.to_string()))
}

pub fn parse_frame(buf: &mut BytesMut) -> Result<Option<RespFrame>, RespError> {
    if !buf.has_remaining() {
        return Ok(None);
    }

    match buf[0] {
        b'+' | b'-' | b':' => {
            let line = read_line(buf)?.ok_or(RespError::Incomplete)?;
            let (ty, data) = line
                .split_first()
                .ok_or_else(|| RespError::Protocol("empty line".into()))?;
            let s =
                std::str::from_utf8(data).map_err(|e| RespError::Protocol(e.to_string()))?;
            let frame = match ty {
                b'+' => RespFrame::SimpleString(s.to_string()),
                b'-' => RespFrame::Error(s.to_string()),
                b':' => RespFrame::Integer(
                    s.parse::<i64>()
                        .map_err(|e| RespError::Protocol(e.to_string()))?,
                ),
                _ => return Err(RespError::Protocol("invalid prefix".into())),
            };
            Ok(Some(frame))
        }
        b'$' => {
            let header = read_line(buf)?.ok_or(RespError::Incomplete)?;
            let len = parse_isize(&header[1..])?;
            if len < 0 {
                return Ok(Some(RespFrame::BulkString(None)));
            }
            let len = len as usize;
            if buf.len() < len + 2 {
                return Err(RespError::Incomplete);
            }
            let data = buf.split_to(len).to_vec();
            if buf.len() < 2 || &buf[..2] != b"\r\n" {
                return Err(RespError::Protocol("bulkstring missing CRLF".into()));
            }
            buf.advance(2);
            Ok(Some(RespFrame::BulkString(Some(data))))
        }
        b'*' => {
            let header = read_line(buf)?.ok_or(RespError::Incomplete)?;
            let len = parse_isize(&header[1..])?;
            if len < 0 {
                return Ok(Some(RespFrame::Array(None)));
            }
            let len = len as usize;
            let mut items = Vec::with_capacity(len);
            for _ in 0..len {
                let frame = parse_frame(buf)?.ok_or(RespError::Incomplete)?;
                items.push(frame);
            }
            Ok(Some(RespFrame::Array(Some(items))))
        }
        _ => Err(RespError::Protocol("unknown RESP type".into())),
    }
}

pub fn encode_frame(frame: &RespFrame) -> Vec<u8> {
    match frame {
        RespFrame::SimpleString(s) => {
            let mut out = Vec::with_capacity(s.len() + 3);
            out.extend_from_slice(b"+");
            out.extend_from_slice(s.as_bytes());
            out.extend_from_slice(b"\r\n");
            out
        }
        RespFrame::Error(s) => {
            let mut out = Vec::with_capacity(s.len() + 3);
            out.extend_from_slice(b"-");
            out.extend_from_slice(s.as_bytes());
            out.extend_from_slice(b"\r\n");
            out
        }
        RespFrame::Integer(i) => {
            let s = i.to_string();
            let mut out = Vec::with_capacity(s.len() + 3);
            out.extend_from_slice(b":");
            out.extend_from_slice(s.as_bytes());
            out.extend_from_slice(b"\r\n");
            out
        }
        RespFrame::BulkString(None) => b"$-1\r\n".to_vec(),
        RespFrame::BulkString(Some(b)) => {
            let len = b.len().to_string();
            let mut out = Vec::with_capacity(len.len() + b.len() + 5);
            out.extend_from_slice(b"$");
            out.extend_from_slice(len.as_bytes());
            out.extend_from_slice(b"\r\n");
            out.extend_from_slice(b);
            out.extend_from_slice(b"\r\n");
            out
        }
        RespFrame::Array(None) => b"*-1\r\n".to_vec(),
        RespFrame::Array(Some(arr)) => {
            let len = arr.len().to_string();
            let mut out = Vec::new();
            out.extend_from_slice(b"*");
            out.extend_from_slice(len.as_bytes());
            out.extend_from_slice(b"\r\n");
            for f in arr {
                out.extend_from_slice(&encode_frame(f));
            }
            out
        }
    }
}