1use std::{convert::TryInto, fmt};
2
3use bytes::{Buf, BufMut};
4use thiserror::Error;
5
6use crate::coding::{self, Codec, UnexpectedEnd};
7
8#[cfg(feature = "arbitrary")]
9use arbitrary::Arbitrary;
10
11#[derive(Default, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
17pub struct VarInt(pub(crate) u64);
18
19impl VarInt {
20 pub const MAX: Self = Self((1 << 62) - 1);
22 pub const MAX_SIZE: usize = 8;
24
25 pub const fn from_u32(x: u32) -> Self {
27 Self(x as u64)
28 }
29
30 pub fn from_u64(x: u64) -> Result<Self, VarIntBoundsExceeded> {
32 if x < 2u64.pow(62) {
33 Ok(Self(x))
34 } else {
35 Err(VarIntBoundsExceeded)
36 }
37 }
38
39 pub const unsafe fn from_u64_unchecked(x: u64) -> Self {
45 Self(x)
46 }
47
48 pub const fn into_inner(self) -> u64 {
50 self.0
51 }
52
53 pub(crate) const fn size(self) -> usize {
55 let x = self.0;
56 if x < 2u64.pow(6) {
57 1
58 } else if x < 2u64.pow(14) {
59 2
60 } else if x < 2u64.pow(30) {
61 4
62 } else if x < 2u64.pow(62) {
63 8
64 } else {
65 panic!("malformed VarInt");
66 }
67 }
68}
69
70impl From<VarInt> for u64 {
71 fn from(x: VarInt) -> Self {
72 x.0
73 }
74}
75
76impl From<u8> for VarInt {
77 fn from(x: u8) -> Self {
78 Self(x.into())
79 }
80}
81
82impl From<u16> for VarInt {
83 fn from(x: u16) -> Self {
84 Self(x.into())
85 }
86}
87
88impl From<u32> for VarInt {
89 fn from(x: u32) -> Self {
90 Self(x.into())
91 }
92}
93
94impl std::convert::TryFrom<u64> for VarInt {
95 type Error = VarIntBoundsExceeded;
96 fn try_from(x: u64) -> Result<Self, VarIntBoundsExceeded> {
98 Self::from_u64(x)
99 }
100}
101
102impl std::convert::TryFrom<u128> for VarInt {
103 type Error = VarIntBoundsExceeded;
104 fn try_from(x: u128) -> Result<Self, VarIntBoundsExceeded> {
106 Self::from_u64(x.try_into().map_err(|_| VarIntBoundsExceeded)?)
107 }
108}
109
110impl std::convert::TryFrom<usize> for VarInt {
111 type Error = VarIntBoundsExceeded;
112 fn try_from(x: usize) -> Result<Self, VarIntBoundsExceeded> {
114 Self::try_from(x as u64)
115 }
116}
117
118impl fmt::Debug for VarInt {
119 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
120 self.0.fmt(f)
121 }
122}
123
124impl fmt::Display for VarInt {
125 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
126 self.0.fmt(f)
127 }
128}
129
130#[cfg(feature = "arbitrary")]
131impl<'arbitrary> Arbitrary<'arbitrary> for VarInt {
132 fn arbitrary(u: &mut arbitrary::Unstructured<'arbitrary>) -> arbitrary::Result<Self> {
133 Ok(Self(u.int_in_range(0..=Self::MAX.0)?))
134 }
135}
136
137#[derive(Debug, Copy, Clone, Eq, PartialEq, Error)]
139#[error("value too large for varint encoding")]
140pub struct VarIntBoundsExceeded;
141
142impl Codec for VarInt {
143 fn decode<B: Buf>(r: &mut B) -> coding::Result<Self> {
144 if !r.has_remaining() {
145 return Err(UnexpectedEnd);
146 }
147 let mut buf = [0; 8];
148 buf[0] = r.get_u8();
149 let tag = buf[0] >> 6;
150 buf[0] &= 0b0011_1111;
151 let x = match tag {
152 0b00 => u64::from(buf[0]),
153 0b01 => {
154 if r.remaining() < 1 {
155 return Err(UnexpectedEnd);
156 }
157 r.copy_to_slice(&mut buf[1..2]);
158 u64::from(u16::from_be_bytes(buf[..2].try_into().unwrap()))
159 }
160 0b10 => {
161 if r.remaining() < 3 {
162 return Err(UnexpectedEnd);
163 }
164 r.copy_to_slice(&mut buf[1..4]);
165 u64::from(u32::from_be_bytes(buf[..4].try_into().unwrap()))
166 }
167 0b11 => {
168 if r.remaining() < 7 {
169 return Err(UnexpectedEnd);
170 }
171 r.copy_to_slice(&mut buf[1..8]);
172 u64::from_be_bytes(buf)
173 }
174 _ => unreachable!(),
175 };
176 Ok(Self(x))
177 }
178
179 fn encode<B: BufMut>(&self, w: &mut B) {
180 let x = self.0;
181 if x < 2u64.pow(6) {
182 w.put_u8(x as u8);
183 } else if x < 2u64.pow(14) {
184 w.put_u16((0b01 << 14) | x as u16);
185 } else if x < 2u64.pow(30) {
186 w.put_u32((0b10 << 30) | x as u32);
187 } else if x < 2u64.pow(62) {
188 w.put_u64((0b11 << 62) | x);
189 } else {
190 unreachable!("malformed VarInt")
191 }
192 }
193}