1use std::io::{Read, Write};
2
3use crate::error::{CodecError, CodecKind, InvalidEncodingReason};
4use crate::io::{read_exact_counted, write_all_counted};
5
6pub trait VarIntWrite: Write {
7 #[inline]
8 fn write_varint(&mut self, value: i32) -> Result<(), CodecError> {
9 self.write_varint_with_size(value).map(|_| ())
10 }
11
12 #[inline]
13 fn write_varint_with_size(&mut self, value: i32) -> Result<usize, CodecError> {
14 let mut value = value as u32;
15 let mut bytes_processed = 0;
16
17 loop {
18 let byte = (value & 0x7F) as u8;
19 value >>= 7;
20 let has_next = value != 0;
21 let byte = if has_next { byte | 0x80 } else { byte };
22
23 write_all_counted(self, &[byte], CodecKind::VarInt, bytes_processed)?;
24 bytes_processed += 1;
25
26 if !has_next {
27 return Ok(bytes_processed);
28 }
29 }
30 }
31}
32
33pub trait VarIntRead: Read {
34 #[inline]
35 fn read_varint(&mut self) -> Result<i32, CodecError> {
36 self.read_varint_with_size().map(|(value, _)| value)
37 }
38
39 #[inline]
40 fn read_varint_with_size(&mut self) -> Result<(i32, usize), CodecError> {
41 let mut result = 0u32;
42 let mut shift = 0;
43
44 for i in 0..5 {
45 let mut buf = [0u8; 1];
46 read_exact_counted(self, &mut buf, CodecKind::VarInt, i)?;
47 let byte = buf[0];
48
49 if i == 4 {
50 if (byte & 0x80) != 0 {
51 return Err(CodecError::invalid_encoding(
52 CodecKind::VarInt,
53 i + 1,
54 InvalidEncodingReason::TooLong { max_bytes: 5 },
55 ));
56 }
57 if (byte & !0x0F) != 0 {
58 return Err(CodecError::invalid_encoding(
59 CodecKind::VarInt,
60 i + 1,
61 InvalidEncodingReason::ValueOutOfRange {
62 terminal_byte: byte,
63 allowed_mask: 0x0F,
64 },
65 ));
66 }
67 }
68
69 let value = (byte & 0x7F) as u32;
70 result |= value << shift;
71
72 if (byte & 0x80) == 0 {
73 return Ok((result as i32, i + 1));
74 }
75
76 shift += 7;
77 }
78
79 unreachable!("the fifth VarInt byte always terminates or returns an error")
80 }
81}
82
83impl<R: Read> VarIntRead for R {}
84impl<W: Write> VarIntWrite for W {}