Skip to main content

kafrust_protocol/codec/
encode.rs

1use crate::error::{Error, Result};
2
3#[derive(Debug, Clone, Default)]
4pub struct Encoder {
5    output: Vec<u8>,
6}
7
8impl Encoder {
9    pub fn new() -> Self {
10        Self::default()
11    }
12
13    pub fn into_bytes(self) -> Vec<u8> {
14        self.output
15    }
16
17    pub fn write_i8(&mut self, value: i8) {
18        self.output.push(value as u8);
19    }
20
21    pub fn write_bool(&mut self, value: bool) {
22        self.write_i8(if value { 1 } else { 0 });
23    }
24
25    pub fn write_i16(&mut self, value: i16) {
26        self.output.extend_from_slice(&value.to_be_bytes());
27    }
28
29    pub fn write_i32(&mut self, value: i32) {
30        self.output.extend_from_slice(&value.to_be_bytes());
31    }
32
33    pub fn write_i64(&mut self, value: i64) {
34        self.output.extend_from_slice(&value.to_be_bytes());
35    }
36
37    pub fn write_string(&mut self, value: &str) -> Result<()> {
38        let length = i16::try_from(value.len()).map_err(|_| Error::LengthOverflow("string"))?;
39        self.write_i16(length);
40        self.output.extend_from_slice(value.as_bytes());
41        Ok(())
42    }
43
44    pub fn write_nullable_string(&mut self, value: Option<&str>) -> Result<()> {
45        match value {
46            Some(value) => self.write_string(value),
47            None => {
48                self.write_i16(-1);
49                Ok(())
50            }
51        }
52    }
53
54    pub fn write_bytes(&mut self, value: &[u8]) -> Result<()> {
55        let length = i32::try_from(value.len()).map_err(|_| Error::LengthOverflow("bytes"))?;
56        self.write_i32(length);
57        self.output.extend_from_slice(value);
58        Ok(())
59    }
60
61    pub fn write_nullable_bytes(&mut self, value: Option<&[u8]>) -> Result<()> {
62        match value {
63            Some(value) => self.write_bytes(value),
64            None => {
65                self.write_i32(-1);
66                Ok(())
67            }
68        }
69    }
70
71    pub fn write_unsigned_varint(&mut self, mut value: u32) {
72        loop {
73            let mut byte = (value & 0x7f) as u8;
74            value >>= 7;
75            if value != 0 {
76                byte |= 0x80;
77            }
78            self.output.push(byte);
79            if value == 0 {
80                break;
81            }
82        }
83    }
84
85    pub fn write_varint(&mut self, value: i32) {
86        self.write_unsigned_varint(((value as u32) << 1) ^ ((value >> 31) as u32));
87    }
88
89    pub fn write_varlong(&mut self, value: i64) {
90        let mut value = ((value as u64) << 1) ^ ((value >> 63) as u64);
91        loop {
92            let mut byte = (value & 0x7f) as u8;
93            value >>= 7;
94            if value != 0 {
95                byte |= 0x80;
96            }
97            self.output.push(byte);
98            if value == 0 {
99                break;
100            }
101        }
102    }
103
104    pub fn write_varint_bytes(&mut self, value: &[u8]) -> Result<()> {
105        let length =
106            i32::try_from(value.len()).map_err(|_| Error::LengthOverflow("varint bytes"))?;
107        self.write_varint(length);
108        self.output.extend_from_slice(value);
109        Ok(())
110    }
111
112    pub fn write_varint_nullable_bytes(&mut self, value: Option<&[u8]>) -> Result<()> {
113        match value {
114            Some(value) => self.write_varint_bytes(value),
115            None => {
116                self.write_varint(-1);
117                Ok(())
118            }
119        }
120    }
121
122    pub fn write_compact_string(&mut self, value: &str) -> Result<()> {
123        let length =
124            u32::try_from(value.len()).map_err(|_| Error::LengthOverflow("compact string"))?;
125        let encoded_length = length
126            .checked_add(1)
127            .ok_or(Error::LengthOverflow("compact string"))?;
128        self.write_unsigned_varint(encoded_length);
129        self.output.extend_from_slice(value.as_bytes());
130        Ok(())
131    }
132
133    pub fn write_compact_nullable_string(&mut self, value: Option<&str>) -> Result<()> {
134        match value {
135            Some(value) => self.write_compact_string(value),
136            None => {
137                self.write_unsigned_varint(0);
138                Ok(())
139            }
140        }
141    }
142
143    pub fn write_compact_bytes(&mut self, value: &[u8]) -> Result<()> {
144        let length =
145            u32::try_from(value.len()).map_err(|_| Error::LengthOverflow("compact bytes"))?;
146        let encoded_length = length
147            .checked_add(1)
148            .ok_or(Error::LengthOverflow("compact bytes"))?;
149        self.write_unsigned_varint(encoded_length);
150        self.output.extend_from_slice(value);
151        Ok(())
152    }
153
154    pub fn write_compact_nullable_bytes(&mut self, value: Option<&[u8]>) -> Result<()> {
155        match value {
156            Some(value) => self.write_compact_bytes(value),
157            None => {
158                self.write_unsigned_varint(0);
159                Ok(())
160            }
161        }
162    }
163
164    pub fn write_array<T>(
165        &mut self,
166        values: Option<&[T]>,
167        mut write_item: impl FnMut(&mut Self, &T) -> Result<()>,
168    ) -> Result<()> {
169        match values {
170            Some(values) => {
171                let length =
172                    i32::try_from(values.len()).map_err(|_| Error::LengthOverflow("array"))?;
173                self.write_i32(length);
174                for value in values {
175                    write_item(self, value)?;
176                }
177            }
178            None => self.write_i32(-1),
179        }
180        Ok(())
181    }
182
183    pub fn write_compact_array<T>(
184        &mut self,
185        values: Option<&[T]>,
186        mut write_item: impl FnMut(&mut Self, &T) -> Result<()>,
187    ) -> Result<()> {
188        match values {
189            Some(values) => {
190                let length =
191                    u32::try_from(values.len()).map_err(|_| Error::LengthOverflow("array"))?;
192                let encoded_length = length
193                    .checked_add(1)
194                    .ok_or(Error::LengthOverflow("compact array"))?;
195                self.write_unsigned_varint(encoded_length);
196                for value in values {
197                    write_item(self, value)?;
198                }
199            }
200            None => self.write_unsigned_varint(0),
201        }
202        Ok(())
203    }
204
205    pub fn write_empty_tagged_fields(&mut self) {
206        self.write_unsigned_varint(0);
207    }
208
209    pub fn write_raw(&mut self, bytes: &[u8]) {
210        self.output.extend_from_slice(bytes);
211    }
212}