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    /// Writes a Kafka UUID as its 16-byte network-order representation.
38    pub fn write_uuid(&mut self, value: &[u8; 16]) {
39        self.output.extend_from_slice(value);
40    }
41
42    pub fn write_f64(&mut self, value: f64) {
43        self.output
44            .extend_from_slice(&value.to_bits().to_be_bytes());
45    }
46
47    pub fn write_string(&mut self, value: &str) -> Result<()> {
48        let length = i16::try_from(value.len()).map_err(|_| Error::LengthOverflow("string"))?;
49        self.write_i16(length);
50        self.output.extend_from_slice(value.as_bytes());
51        Ok(())
52    }
53
54    pub fn write_nullable_string(&mut self, value: Option<&str>) -> Result<()> {
55        match value {
56            Some(value) => self.write_string(value),
57            None => {
58                self.write_i16(-1);
59                Ok(())
60            }
61        }
62    }
63
64    pub fn write_bytes(&mut self, value: &[u8]) -> Result<()> {
65        let length = i32::try_from(value.len()).map_err(|_| Error::LengthOverflow("bytes"))?;
66        self.write_i32(length);
67        self.output.extend_from_slice(value);
68        Ok(())
69    }
70
71    pub fn write_nullable_bytes(&mut self, value: Option<&[u8]>) -> Result<()> {
72        match value {
73            Some(value) => self.write_bytes(value),
74            None => {
75                self.write_i32(-1);
76                Ok(())
77            }
78        }
79    }
80
81    pub fn write_unsigned_varint(&mut self, mut value: u32) {
82        loop {
83            let mut byte = (value & 0x7f) as u8;
84            value >>= 7;
85            if value != 0 {
86                byte |= 0x80;
87            }
88            self.output.push(byte);
89            if value == 0 {
90                break;
91            }
92        }
93    }
94
95    pub fn write_varint(&mut self, value: i32) {
96        self.write_unsigned_varint(((value as u32) << 1) ^ ((value >> 31) as u32));
97    }
98
99    pub fn write_varlong(&mut self, value: i64) {
100        let mut value = ((value as u64) << 1) ^ ((value >> 63) as u64);
101        loop {
102            let mut byte = (value & 0x7f) as u8;
103            value >>= 7;
104            if value != 0 {
105                byte |= 0x80;
106            }
107            self.output.push(byte);
108            if value == 0 {
109                break;
110            }
111        }
112    }
113
114    pub fn write_varint_bytes(&mut self, value: &[u8]) -> Result<()> {
115        let length =
116            i32::try_from(value.len()).map_err(|_| Error::LengthOverflow("varint bytes"))?;
117        self.write_varint(length);
118        self.output.extend_from_slice(value);
119        Ok(())
120    }
121
122    pub fn write_varint_nullable_bytes(&mut self, value: Option<&[u8]>) -> Result<()> {
123        match value {
124            Some(value) => self.write_varint_bytes(value),
125            None => {
126                self.write_varint(-1);
127                Ok(())
128            }
129        }
130    }
131
132    pub fn write_compact_string(&mut self, value: &str) -> Result<()> {
133        let length =
134            u32::try_from(value.len()).map_err(|_| Error::LengthOverflow("compact string"))?;
135        let encoded_length = length
136            .checked_add(1)
137            .ok_or(Error::LengthOverflow("compact string"))?;
138        self.write_unsigned_varint(encoded_length);
139        self.output.extend_from_slice(value.as_bytes());
140        Ok(())
141    }
142
143    pub fn write_compact_nullable_string(&mut self, value: Option<&str>) -> Result<()> {
144        match value {
145            Some(value) => self.write_compact_string(value),
146            None => {
147                self.write_unsigned_varint(0);
148                Ok(())
149            }
150        }
151    }
152
153    pub fn write_compact_bytes(&mut self, value: &[u8]) -> Result<()> {
154        let length =
155            u32::try_from(value.len()).map_err(|_| Error::LengthOverflow("compact bytes"))?;
156        let encoded_length = length
157            .checked_add(1)
158            .ok_or(Error::LengthOverflow("compact bytes"))?;
159        self.write_unsigned_varint(encoded_length);
160        self.output.extend_from_slice(value);
161        Ok(())
162    }
163
164    pub fn write_compact_nullable_bytes(&mut self, value: Option<&[u8]>) -> Result<()> {
165        match value {
166            Some(value) => self.write_compact_bytes(value),
167            None => {
168                self.write_unsigned_varint(0);
169                Ok(())
170            }
171        }
172    }
173
174    pub fn write_array<T>(
175        &mut self,
176        values: Option<&[T]>,
177        mut write_item: impl FnMut(&mut Self, &T) -> Result<()>,
178    ) -> Result<()> {
179        match values {
180            Some(values) => {
181                let length =
182                    i32::try_from(values.len()).map_err(|_| Error::LengthOverflow("array"))?;
183                self.write_i32(length);
184                for value in values {
185                    write_item(self, value)?;
186                }
187            }
188            None => self.write_i32(-1),
189        }
190        Ok(())
191    }
192
193    pub fn write_compact_array<T>(
194        &mut self,
195        values: Option<&[T]>,
196        mut write_item: impl FnMut(&mut Self, &T) -> Result<()>,
197    ) -> Result<()> {
198        match values {
199            Some(values) => {
200                let length =
201                    u32::try_from(values.len()).map_err(|_| Error::LengthOverflow("array"))?;
202                let encoded_length = length
203                    .checked_add(1)
204                    .ok_or(Error::LengthOverflow("compact array"))?;
205                self.write_unsigned_varint(encoded_length);
206                for value in values {
207                    write_item(self, value)?;
208                }
209            }
210            None => self.write_unsigned_varint(0),
211        }
212        Ok(())
213    }
214
215    pub fn write_empty_tagged_fields(&mut self) {
216        self.write_unsigned_varint(0);
217    }
218
219    pub fn write_raw(&mut self, bytes: &[u8]) {
220        self.output.extend_from_slice(bytes);
221    }
222}
kafrust_protocol/codec/encode.rs

kafrust_protocol/codec/
encode.rs
