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
7pub 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 pub fn new(buf: &'a mut [u8]) -> Self {
17 VarIntEncoder {
18 buf,
19 pos: 0,
20 _marker: PhantomData,
21 }
22 }
23
24 pub fn position(&self) -> usize {
26 self.pos
27 }
28
29 pub fn remaining(&self) -> usize {
31 self.buf.len() - self.pos
32 }
33
34 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 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 #[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 #[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 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 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
113pub 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 pub fn new(buf: &'a [u8]) -> Self {
123 VarIntDecoder {
124 buf,
125 pos: 0,
126 _marker: PhantomData,
127 }
128 }
129
130 pub fn position(&self) -> usize {
132 self.pos
133 }
134
135 pub fn remaining(&self) -> &'a [u8] {
137 &self.buf[self.pos..]
138 }
139
140 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 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 #[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 #[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 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 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#[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#[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}