tiny_varint/
batch.rs

1use crate::error::Error;
2use crate::traits::VarInt;
3use crate::encoding::{encode, decode};
4use crate::zigzag::{ZigZag, encode_zigzag, decode_zigzag};
5use core::marker::PhantomData;
6
7/// Batch encoder for VarInt values with state management
8pub struct VarIntEncoder<'a, T: VarInt = u64> {
9    buf: &'a mut [u8],
10    pos: usize,
11    _marker: PhantomData<T>,
12}
13
14impl<'a, T: VarInt> VarIntEncoder<'a, T> {
15    /// Creates a new encoder with the provided buffer
16    pub fn new(buf: &'a mut [u8]) -> Self {
17        VarIntEncoder {
18            buf,
19            pos: 0,
20            _marker: PhantomData,
21        }
22    }
23    
24    /// Gets the current position in the buffer
25    pub fn position(&self) -> usize {
26        self.pos
27    }
28    
29    /// Gets the remaining space in the buffer
30    pub fn remaining(&self) -> usize {
31        self.buf.len() - self.pos
32    }
33    
34    /// Writes a VarInt value to the buffer
35    ///
36    /// Returns the number of bytes written
37    pub fn write(&mut self, value: T) -> Result<usize, Error> {
38        if self.pos >= self.buf.len() {
39            return Err(Error::BufferTooSmall {
40                needed: self.pos + 1,
41                actual: self.buf.len(),
42            });
43        }
44        
45        let bytes_written = encode(value, &mut self.buf[self.pos..])?;
46        self.pos += bytes_written;
47        Ok(bytes_written)
48    }
49    
50    /// Writes a batch of VarInt values
51    ///
52    /// Returns the total number of bytes written
53    pub fn write_batch(&mut self, values: &[T]) -> Result<usize, Error> {
54        let start_pos = self.pos;
55        for &value in values {
56            self.write(value)?;
57        }
58        Ok(self.pos - start_pos)
59    }
60    
61    /// Writes a u64 value to the buffer (convenience method)
62    ///
63    /// Returns the number of bytes written
64    #[inline]
65    pub fn write_u64(&mut self, value: u64) -> Result<usize, Error> 
66    where T: From<u64> {
67        self.write(T::from(value))
68    }
69    
70    /// Writes a u128 value to the buffer (convenience method)
71    ///
72    /// Returns the number of bytes written
73    #[inline]
74    pub fn write_u128(&mut self, value: u128) -> Result<usize, Error>
75    where T: From<u128> {
76        self.write(T::from(value))
77    }
78    
79    /// Writes a signed value using zigzag encoding
80    ///
81    /// Returns the number of bytes written
82    pub fn write_zigzag<S>(&mut self, value: S) -> Result<usize, Error>
83    where 
84        S: ZigZag,
85        S::Unsigned: VarInt {
86        if self.pos >= self.buf.len() {
87            return Err(Error::BufferTooSmall {
88                needed: self.pos + 1,
89                actual: self.buf.len(),
90            });
91        }
92        
93        let bytes_written = encode_zigzag(value, &mut self.buf[self.pos..])?;
94        self.pos += bytes_written;
95        Ok(bytes_written)
96    }
97    
98    /// Writes a batch of signed values using zigzag encoding
99    ///
100    /// Returns the total number of bytes written
101    pub fn write_zigzag_batch<S>(&mut self, values: &[S]) -> Result<usize, Error>
102    where 
103        S: ZigZag,
104        S::Unsigned: VarInt {
105        let start_pos = self.pos;
106        for &value in values {
107            self.write_zigzag(value)?;
108        }
109        Ok(self.pos - start_pos)
110    }
111}
112
113/// Batch decoder for VarInt values with state management
114pub struct VarIntDecoder<'a, T: VarInt = u64> {
115    buf: &'a [u8],
116    pos: usize,
117    _marker: PhantomData<T>,
118}
119
120impl<'a, T: VarInt> VarIntDecoder<'a, T> {
121    /// Creates a new decoder with the provided buffer
122    pub fn new(buf: &'a [u8]) -> Self {
123        VarIntDecoder {
124            buf,
125            pos: 0,
126            _marker: PhantomData,
127        }
128    }
129    
130    /// Gets the current position in the buffer
131    pub fn position(&self) -> usize {
132        self.pos
133    }
134    
135    /// Gets the remaining bytes in the buffer
136    pub fn remaining(&self) -> &'a [u8] {
137        &self.buf[self.pos..]
138    }
139    
140    /// Reads a VarInt value from the buffer
141    ///
142    /// Returns the decoded value
143    pub fn read(&mut self) -> Result<T, Error> {
144        if self.pos >= self.buf.len() {
145            return Err(Error::InputTooShort);
146        }
147        
148        let (value, bytes_read) = decode(&self.buf[self.pos..])?;
149        self.pos += bytes_read;
150        Ok(value)
151    }
152    
153    /// Reads a batch of VarInt values into the provided buffer
154    ///
155    /// Returns the number of values read
156    pub fn read_batch(&mut self, values: &mut [T]) -> Result<usize, Error> {
157        let mut count = 0;
158        
159        while count < values.len() && self.pos < self.buf.len() {
160            match self.read() {
161                Ok(value) => {
162                    values[count] = value;
163                    count += 1;
164                }
165                Err(Error::InputTooShort) => break,
166                Err(e) => return Err(e),
167            }
168        }
169        
170        Ok(count)
171    }
172    
173    /// Reads a u64 value from the buffer (convenience method)
174    ///
175    /// Returns the decoded value
176    #[inline]
177    pub fn read_u64(&mut self) -> Result<u64, Error> 
178    where u64: From<T> {
179        Ok(u64::from(self.read()?))
180    }
181    
182    /// Reads a u128 value from the buffer (convenience method)
183    ///
184    /// Returns the decoded value
185    #[inline]
186    pub fn read_u128(&mut self) -> Result<u128, Error>
187    where u128: From<T> {
188        Ok(u128::from(self.read()?))
189    }
190    
191    /// Reads a signed value that was encoded using zigzag encoding
192    ///
193    /// Returns the decoded value
194    pub fn read_zigzag<S>(&mut self) -> Result<S, Error>
195    where 
196        S: ZigZag,
197        S::Unsigned: VarInt {
198        if self.pos >= self.buf.len() {
199            return Err(Error::InputTooShort);
200        }
201        
202        let (value, bytes_read) = decode_zigzag(&self.buf[self.pos..])?;
203        self.pos += bytes_read;
204        Ok(value)
205    }
206    
207    /// Reads a batch of signed values that were encoded using zigzag encoding
208    ///
209    /// Returns the number of values read
210    pub fn read_zigzag_batch<S>(&mut self, values: &mut [S]) -> Result<usize, Error>
211    where 
212        S: ZigZag,
213        S::Unsigned: VarInt {
214        let mut count = 0;
215        
216        while count < values.len() && self.pos < self.buf.len() {
217            match self.read_zigzag() {
218                Ok(value) => {
219                    values[count] = value;
220                    count += 1;
221                }
222                Err(Error::InputTooShort) => break,
223                Err(e) => return Err(e),
224            }
225        }
226        
227        Ok(count)
228    }
229}
230
231/// Convenience function to encode a batch of u64 values
232///
233/// Returns the number of bytes written
234#[inline]
235pub fn encode_batch(values: &[u64], buf: &mut [u8]) -> Result<usize, Error> {
236    let mut encoder = VarIntEncoder::new(buf);
237    encoder.write_batch(values)
238}
239
240/// Convenience function to decode a batch of u64 values
241///
242/// Returns the number of values read
243#[inline]
244pub fn decode_batch(buf: &[u8], values: &mut [u64]) -> Result<usize, Error> {
245    let mut decoder = VarIntDecoder::new(buf);
246    decoder.read_batch(values)
247}