Skip to main content

ember_protocol/
serialize.rs

1//! Direct-to-buffer RESP3 serialization.
2//!
3//! Writes frames directly into a `BytesMut` buffer with no intermediate
4//! allocations. Integer-to-string conversion uses `itoa` for fast
5//! stack-based formatting.
6
7use bytes::BufMut;
8use bytes::BytesMut;
9
10use crate::types::Frame;
11
12/// Maximum nesting depth for serialization. Matches the parser's limit.
13const MAX_SERIALIZE_DEPTH: usize = 64;
14
15impl Frame {
16    /// Serializes this frame into the provided buffer.
17    ///
18    /// Writes the full RESP3 wire representation, including type prefix
19    /// and trailing `\r\n` delimiters. Nested structures are bounded to
20    /// [`MAX_SERIALIZE_DEPTH`] levels to prevent stack overflow.
21    #[inline]
22    pub fn serialize(&self, dst: &mut BytesMut) {
23        self.serialize_inner(dst, 0);
24    }
25
26    fn serialize_inner(&self, dst: &mut BytesMut, depth: usize) {
27        use crate::types::wire;
28
29        match self {
30            Frame::Simple(s) => match s.as_str() {
31                "OK" => dst.put_slice(wire::OK),
32                "PONG" => dst.put_slice(wire::PONG),
33                _ => {
34                    dst.put_u8(b'+');
35                    dst.put_slice(s.as_bytes());
36                    dst.put_slice(b"\r\n");
37                }
38            },
39            Frame::Error(msg) => {
40                dst.put_u8(b'-');
41                dst.put_slice(msg.as_bytes());
42                dst.put_slice(b"\r\n");
43            }
44            Frame::Integer(n) => match *n {
45                0 => dst.put_slice(wire::ZERO),
46                1 => dst.put_slice(wire::ONE),
47                -1 => dst.put_slice(wire::NEG_ONE),
48                _ => {
49                    dst.put_u8(b':');
50                    write_i64(*n, dst);
51                    dst.put_slice(b"\r\n");
52                }
53            },
54            Frame::Bulk(data) => {
55                dst.put_u8(b'$');
56                write_i64(data.len() as i64, dst);
57                dst.put_slice(b"\r\n");
58                dst.put_slice(data);
59                dst.put_slice(b"\r\n");
60            }
61            Frame::Array(items) => {
62                if depth >= MAX_SERIALIZE_DEPTH {
63                    // emit an error frame instead of overflowing the stack
64                    dst.put_slice(b"-ERR nesting too deep\r\n");
65                    return;
66                }
67                dst.put_u8(b'*');
68                write_i64(items.len() as i64, dst);
69                dst.put_slice(b"\r\n");
70                for item in items {
71                    item.serialize_inner(dst, depth + 1);
72                }
73            }
74            Frame::Null => {
75                dst.put_slice(wire::NULL);
76            }
77            Frame::Map(pairs) => {
78                if depth >= MAX_SERIALIZE_DEPTH {
79                    dst.put_slice(b"-ERR nesting too deep\r\n");
80                    return;
81                }
82                dst.put_u8(b'%');
83                write_i64(pairs.len() as i64, dst);
84                dst.put_slice(b"\r\n");
85                for (key, val) in pairs {
86                    key.serialize_inner(dst, depth + 1);
87                    val.serialize_inner(dst, depth + 1);
88                }
89            }
90        }
91    }
92}
93
94/// Writes an i64 as its decimal ASCII representation directly into the buffer.
95#[inline]
96fn write_i64(val: i64, dst: &mut BytesMut) {
97    let mut buf = itoa::Buffer::new();
98    dst.put_slice(buf.format(val).as_bytes());
99}
100
101#[cfg(test)]
102mod tests {
103    use bytes::Bytes;
104
105    use super::*;
106
107    fn serialize(frame: &Frame) -> Vec<u8> {
108        let mut buf = BytesMut::new();
109        frame.serialize(&mut buf);
110        buf.to_vec()
111    }
112
113    #[test]
114    fn simple_string() {
115        assert_eq!(serialize(&Frame::Simple("OK".into())), b"+OK\r\n");
116    }
117
118    #[test]
119    fn error() {
120        assert_eq!(serialize(&Frame::Error("ERR bad".into())), b"-ERR bad\r\n");
121    }
122
123    #[test]
124    fn integer() {
125        assert_eq!(serialize(&Frame::Integer(42)), b":42\r\n");
126        assert_eq!(serialize(&Frame::Integer(-1)), b":-1\r\n");
127        assert_eq!(serialize(&Frame::Integer(0)), b":0\r\n");
128    }
129
130    #[test]
131    fn bulk_string() {
132        assert_eq!(
133            serialize(&Frame::Bulk(Bytes::from_static(b"hello"))),
134            b"$5\r\nhello\r\n"
135        );
136    }
137
138    #[test]
139    fn empty_bulk_string() {
140        assert_eq!(
141            serialize(&Frame::Bulk(Bytes::from_static(b""))),
142            b"$0\r\n\r\n"
143        );
144    }
145
146    #[test]
147    fn null() {
148        assert_eq!(serialize(&Frame::Null), b"_\r\n");
149    }
150
151    #[test]
152    fn array() {
153        let frame = Frame::Array(vec![Frame::Simple("hello".into()), Frame::Integer(42)]);
154        assert_eq!(serialize(&frame), b"*2\r\n+hello\r\n:42\r\n");
155    }
156
157    #[test]
158    fn empty_array() {
159        assert_eq!(serialize(&Frame::Array(vec![])), b"*0\r\n");
160    }
161
162    #[test]
163    fn map() {
164        let frame = Frame::Map(vec![(Frame::Simple("key".into()), Frame::Integer(1))]);
165        assert_eq!(serialize(&frame), b"%1\r\n+key\r\n:1\r\n");
166    }
167
168    #[test]
169    fn round_trip() {
170        use crate::parse::parse_frame;
171
172        let frames = vec![
173            Frame::Simple("OK".into()),
174            Frame::Error("ERR nope".into()),
175            Frame::Integer(i64::MAX),
176            Frame::Integer(i64::MIN),
177            Frame::Bulk(Bytes::from_static(b"binary\x00data")),
178            Frame::Bulk(Bytes::from_static(b"")),
179            Frame::Null,
180            Frame::Array(vec![
181                Frame::Integer(1),
182                Frame::Bulk(Bytes::from_static(b"two")),
183                Frame::Null,
184            ]),
185            Frame::Map(vec![
186                (Frame::Simple("a".into()), Frame::Integer(1)),
187                (Frame::Simple("b".into()), Frame::Integer(2)),
188            ]),
189            Frame::Array(vec![
190                Frame::Array(vec![Frame::Integer(1), Frame::Integer(2)]),
191                Frame::Array(vec![Frame::Integer(3)]),
192            ]),
193        ];
194
195        for original in &frames {
196            let mut buf = BytesMut::new();
197            original.serialize(&mut buf);
198
199            let (parsed, consumed) = parse_frame(&buf)
200                .expect("round-trip parse should not error")
201                .expect("round-trip parse should return a frame");
202
203            assert_eq!(&parsed, original, "round-trip failed for {original:?}");
204            assert_eq!(consumed, buf.len(), "should consume entire buffer");
205        }
206    }
207}