1use std::io::{self, Read, Write};
40
41#[inline]
42pub(crate) fn write_u8_to<W: Write + ?Sized>(n: u8, writer: &mut W) -> io::Result<usize> {
43 writer.write_all(&[n])?;
44
45 Ok(1)
46}
47
48#[inline]
49pub(crate) fn read_u8_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<u8> {
50 let mut result = [0];
51
52 reader.read_exact(&mut result)?;
53
54 Ok(u8::from_le_bytes(result))
55}
56
57#[inline]
58pub(crate) fn write_i8_to<W: Write + ?Sized>(n: i8, writer: &mut W) -> io::Result<usize> {
59 writer.write_all(&n.to_le_bytes())?;
60
61 Ok(1)
62}
63
64#[inline]
65pub(crate) fn read_i8_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<i8> {
66 let mut result = [0];
67
68 reader.read_exact(&mut result)?;
69
70 Ok(result[0].cast_signed())
71}
72
73#[inline]
74#[allow(clippy::cast_possible_truncation, reason = "False positive.")]
75pub(crate) fn write_u128_to<W: Write + ?Sized>(mut n: u128, writer: &mut W) -> io::Result<usize> {
76 let mut buf = [0u8; 19];
77 let mut idx = 0;
78
79 while n >= 0x80 {
80 buf[idx] = (n as u8) | 0x80;
81 n >>= 7;
82 idx += 1;
83 }
84
85 buf[idx] = n as u8;
86 idx += 1;
87
88 writer.write_all(&buf[..idx])?;
89 Ok(idx)
90}
91
92#[inline]
93#[allow(clippy::cast_lossless, reason = "False positive.")]
94pub(crate) fn read_u128_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<u128> {
95 let mut result = 0u128;
96 let mut shift = 0;
97 let mut buf = [0u8; 1];
98
99 loop {
100 reader.read_exact(&mut buf)?;
101 let byte = buf[0];
102
103 if shift >= 126 && (byte & 0x80) != 0 {
104 return Err(io::Error::new(
105 io::ErrorKind::InvalidData,
106 "varint too large for u128",
107 ));
108 }
109
110 result |= ((byte & 0x7f) as u128) << shift;
111
112 if byte & 0x80 == 0 {
113 break;
114 }
115
116 shift += 7;
117 }
118
119 Ok(result)
120}
121
122macro_rules! impl_unsigned {
123 ($ty:ty) => {
124 paste::paste! {
125 #[inline]
126 #[allow(clippy::cast_lossless, reason = "It is generated code.")]
127 #[allow(clippy::cast_possible_truncation, reason = "It is generated code.")]
128 pub(crate) fn [<write_ $ty _to>]<W: Write + ?Sized>(
129 value: $ty,
130 writer: &mut W,
131 ) -> io::Result<usize> {
132 write_u128_to(value as u128, writer)
133 }
134
135 #[inline]
136 #[allow(clippy::cast_lossless, reason = "It is generated code.")]
137 #[allow(clippy::cast_possible_truncation, reason = "It is generated code.")]
138 pub(crate) fn [<read_ $ty _from>]<R: Read + ?Sized>(
139 reader: &mut R,
140 ) -> io::Result<$ty> {
141 let value = read_u128_from(reader)?;
142
143 if value > <$ty>::MAX as u128 {
144 return Err(io::Error::new(
145 io::ErrorKind::InvalidData,
146 concat!("varint too large for ", stringify!($ty)),
147 ));
148 }
149
150 Ok(value as $ty)
151 }
152
153 pub trait [<Write $ty:camel>]: Write {
154 fn [<write_ $ty>](
155 &mut self,
156 value: $ty,
157 ) -> io::Result<usize> {
158 [<write_ $ty _to>](value, self)
159 }
160 }
161
162 impl<T: Write> [<Write $ty:camel>] for T {}
163
164 pub trait [<Read $ty:camel>]: Read {
165 fn [<read_ $ty>](
166 &mut self,
167 ) -> io::Result<$ty> {
168 [<read_ $ty _from>](self)
169 }
170 }
171
172 impl<T: Read> [<Read $ty:camel>] for T {}
173 }
174 };
175}
176
177impl_unsigned!(u16);
178impl_unsigned!(u32);
179impl_unsigned!(u64);
180
181#[inline]
182#[allow(clippy::cast_sign_loss, reason = "It will be restored")]
183fn zigzag_encode(v: i128) -> u128 {
184 ((v << 1) ^ (v >> 127)) as u128
185}
186
187#[inline]
188#[allow(clippy::cast_possible_wrap, reason = "It will be restored")]
189fn zigzag_decode(v: u128) -> i128 {
190 ((v >> 1) as i128) ^ (-((v & 1) as i128))
191}
192
193#[inline]
194pub(crate) fn write_i128_to<W: Write + ?Sized>(n: i128, writer: &mut W) -> io::Result<usize> {
195 write_u128_to(zigzag_encode(n), writer)
196}
197
198#[inline]
199pub(crate) fn read_i128_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<i128> {
200 Ok(zigzag_decode(read_u128_from(reader)?))
201}
202
203macro_rules! impl_signed {
204 ($ty:ty) => {
205 paste::paste! {
206 #[inline]
207 #[allow(clippy::cast_lossless, reason = "It is generated code.")]
208 #[allow(clippy::cast_possible_truncation, reason = "It is generated code.")]
209 pub(crate) fn [<write_ $ty _to>]<W: Write + ?Sized>(
210 value: $ty,
211 writer: &mut W,
212 ) -> io::Result<usize> {
213 write_u128_to(zigzag_encode(value as i128), writer)
214 }
215
216 #[inline]
217 #[allow(clippy::cast_lossless, reason = "It is generated code.")]
218 #[allow(clippy::cast_possible_truncation, reason = "It is generated code.")]
219 pub(crate) fn [<read_ $ty _from>]<R: Read + ?Sized>(
220 reader: &mut R,
221 ) -> io::Result<$ty> {
222 let value = zigzag_decode(read_u128_from(reader)?);
223
224 if value < <$ty>::MIN as i128 || value > <$ty>::MAX as i128 {
225 return Err(io::Error::new(
226 io::ErrorKind::InvalidData,
227 concat!("varint too large for ", stringify!($ty)),
228 ));
229 }
230
231 Ok(value as $ty)
232 }
233 }
234 };
235}
236
237impl_signed!(i16);
238impl_signed!(i32);
239impl_signed!(i64);
240
241pub trait VarInt: Sized {
251 fn write_as_varint_to<W: Write + ?Sized>(self, writer: &mut W) -> io::Result<usize>;
253 fn read_varint_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<Self>;
255}
256
257macro_rules! impl_varint {
258 ($($ty:ty),*) => {
259 $(
260 paste::paste! {
261 impl VarInt for $ty {
262 #[inline]
263 fn write_as_varint_to<W: Write + ?Sized>(self, writer: &mut W) -> io::Result<usize> {
264 [<write_ $ty _to>](self, writer)
265 }
266
267 #[inline]
268 fn read_varint_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<Self> {
269 [<read_ $ty _from>](reader)
270 }
271 }
272 }
273 )*
274 };
275}
276
277impl_varint!(u8, u16, u32, u64, u128, i8, i16, i32, i64, i128);
278
279pub trait WriteVarInt: Write {
294 fn write_varint<T: VarInt>(&mut self, value: T) -> io::Result<usize> {
307 value.write_as_varint_to(self)
308 }
309}
310
311pub trait ReadVarInt: Read {
331 fn read_varint<T: VarInt>(&mut self) -> io::Result<T> {
349 T::read_varint_from(self)
350 }
351}
352
353impl<W: Write> WriteVarInt for W {}
354impl<R: Read> ReadVarInt for R {}
355
356#[cfg(test)]
357mod tests {
358 use super::*;
359 use std::io::Cursor;
360
361 fn roundtrip<T>(value: T)
362 where
363 T: VarInt + Copy + PartialEq + std::fmt::Debug,
364 {
365 let mut buf = Vec::new();
366
367 value.write_as_varint_to(&mut buf).unwrap();
368
369 let decoded = T::read_varint_from(&mut Cursor::new(buf)).unwrap();
370
371 assert_eq!(value, decoded);
372 }
373
374 #[test]
375 fn roundtrip_unsigned() {
376 roundtrip(0u8);
377 roundtrip(1u8);
378 roundtrip(u8::MAX);
379
380 roundtrip(0u16);
381 roundtrip(127u16);
382 roundtrip(128u16);
383 roundtrip(u16::MAX);
384
385 roundtrip(0u32);
386 roundtrip(u32::MAX);
387
388 roundtrip(0u64);
389 roundtrip(u64::MAX);
390
391 roundtrip(0u128);
392 roundtrip(u128::MAX);
393 }
394
395 #[test]
396 fn roundtrip_signed() {
397 roundtrip(0i8);
398 roundtrip(-1i8);
399 roundtrip(i8::MIN);
400 roundtrip(i8::MAX);
401
402 roundtrip(0i16);
403 roundtrip(-1i16);
404 roundtrip(i16::MIN);
405 roundtrip(i16::MAX);
406
407 roundtrip(0i32);
408 roundtrip(i32::MIN);
409 roundtrip(i32::MAX);
410
411 roundtrip(0i64);
412 roundtrip(i64::MIN);
413 roundtrip(i64::MAX);
414
415 roundtrip(0i128);
416 roundtrip(i128::MIN);
417 roundtrip(i128::MAX);
418 }
419
420 #[test]
421 fn encoding_size() {
422 let mut buf = Vec::new();
423
424 assert_eq!(write_u128_to(0, &mut buf).unwrap(), 1);
425
426 buf.clear();
427 assert_eq!(write_u128_to(127, &mut buf).unwrap(), 1);
428
429 buf.clear();
430 assert_eq!(write_u128_to(128, &mut buf).unwrap(), 2);
431
432 buf.clear();
433 assert_eq!(write_u128_to(16383, &mut buf).unwrap(), 2);
434
435 buf.clear();
436 assert_eq!(write_u128_to(16384, &mut buf).unwrap(), 3);
437 }
438
439 #[test]
440 fn malformed_varint_is_rejected() {
441 let bytes = [0xff; 19];
442
443 let err = read_u128_from(&mut Cursor::new(bytes)).unwrap_err();
444
445 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
446 }
447
448 #[test]
449 fn overflow_for_smaller_type() {
450 let mut buf = Vec::new();
451
452 write_u32_to(70000, &mut buf).unwrap();
453
454 let err = read_u16_from(&mut Cursor::new(buf)).unwrap_err();
455
456 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
457 }
458
459 #[test]
460 fn generic_traits_work() {
461 let mut buf = Vec::new();
462
463 buf.write_varint(12345u32).unwrap();
464 buf.write_varint(-567i32).unwrap();
465
466 let mut cursor = Cursor::new(buf);
467
468 let a: u32 = cursor.read_varint().unwrap();
469 let b: i32 = cursor.read_varint().unwrap();
470
471 assert_eq!(a, 12345);
472 assert_eq!(b, -567);
473 }
474}