varint.rs

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct VarInt(u128);

macro_rules! conversions {
	($($ty:ty => $unsigned:ty, $signed:expr);*;) => {$(
		impl From<$ty> for VarInt {
			fn from(v: $ty) -> Self {
				let raw = if $signed {
					// rotate sign bit to front, and flip value bits if negative
					// this encodes small absolute values in the fewest bytes possible
					let raw = (v as $unsigned).rotate_left(1);
					#[allow(unused_comparisons)]
					if v < 0 {
						raw ^ !1
					} else {
                        raw
                    }
				} else {
                    v as _
                };
				Self(raw as _)
			}
		}

		impl From<VarInt> for $ty {
			fn from(VarInt(raw): VarInt) -> Self {
				if $signed {
					let raw = if raw & 1 > 0 {
						raw ^ !1
					} else {
                        raw
                    };
					raw.rotate_right(1) as $ty
				} else {
                    raw as _
                }
			}
		}
	)*}
}

conversions! {
	u8 => u8, false;
	i8 => u8, true;

	u16 => u16, false;
	i16 => u16, true;

	u32 => u32, false;
	i32 => u32, true;

	u64 => u64, false;
	i64 => u64, true;

	u128 => u128, false;
	i128 => u128, true;

	usize => usize, false;
	isize => usize, true;
}

#[test]
fn test_varint_conversions() {
	for v in 0 ..= u8::MAX {
		let vi = VarInt::from(v);
		let vp = vi.into();
		assert_eq!(v, vp);
	}

	for v in i8::MIN ..= i8::MAX {
		let vi = VarInt::from(v);
		let vp = vi.into();
		assert_eq!(v, vp);
	}

	for v in u16::MIN ..= u16::MAX {
		let vi = VarInt::from(v);
		let vp = vi.into();
		assert_eq!(v, vp);
	}

	for v in i16::MIN ..= i16::MAX {
		let vi = VarInt::from(v);
		let vp = vi.into();
		assert_eq!(v, vp);
	}
}

#[derive(Clone, Copy, Debug)]
pub enum VarIntError {
	NoData,
	IncompleteData,
	MalformedData,
}

impl VarInt {
	const MAX_BYTES: usize = ((u128::BITS + 6) / 7) as usize;

	pub fn from_bytes(bytes: &[u8]) -> Result<(Self, usize), VarIntError> {
		if bytes.is_empty() {
			return Err(VarIntError::NoData);
		}

		let mut len = 0;
		loop {
			if len > Self::MAX_BYTES {
				return Err(VarIntError::MalformedData);
			}
			if len >= bytes.len() {
				return Err(VarIntError::IncompleteData);
			}

			if bytes[len] & 128 == 0 {
				break;
			}
			len += 1;
		}

		let mut raw = 0u128;
		for i in (0 ..= len).rev() {
			raw <<= 7;
			raw |= (bytes[i] & 127) as u128;
		}
		Ok((Self(raw), len + 1))
	}

	pub fn to_bytes(self) -> impl Iterator<Item = u8> {
		let mut raw = self.0;
		let mut res = [0u8; Self::MAX_BYTES];
		if raw == 0 {
			return res.into_iter().take(1);
		}

		let mut len = 0;
		while raw > 0 {
			res[len] = (raw & 127) as u8;
			raw >>= 7;
			if raw > 0 {
				res[len] |= 128;
			}
			len += 1;
		}

		res.into_iter().take(len)
	}
}

#[test]
fn test_varint_codec() {
	for v in 0 .. 128 {
		let bytes: Vec<_> = VarInt::from(v).to_bytes().collect();
		assert_eq!(bytes.len(), 1);
		assert_eq!(bytes[0], v);
	}

	for v in 128 ..= u16::MAX {
		let bytes: Vec<_> = VarInt::from(v).to_bytes().collect();
		let (vp, rlen) = VarInt::from_bytes(&bytes).unwrap();
		assert_eq!(bytes.len(), rlen);
		assert_eq!(v, u16::from(vp));
	}
}