use crate::buf::{ReadBuf, WriteBuf};
use crate::error::{Error, Result};
use crate::zigzag;
pub const MAX_LEN_U16: usize = 3;
pub const MAX_LEN_U32: usize = 5;
pub const MAX_LEN_U64: usize = 10;
#[inline]
#[must_use]
pub const fn encoded_len_u16(value: u16) -> usize {
encoded_len_u64(value as u64)
}
#[inline]
#[must_use]
pub const fn encoded_len_u32(value: u32) -> usize {
encoded_len_u64(value as u64)
}
#[inline]
#[must_use]
pub const fn encoded_len_u64(value: u64) -> usize {
let bits = (u64::BITS - value.leading_zeros()) as usize;
if bits == 0 { 1 } else { bits.div_ceil(7) }
}
#[inline]
#[must_use]
pub const fn encoded_len_usize(value: usize) -> usize {
encoded_len_u64(value as u64)
}
#[inline]
#[must_use]
pub const fn encoded_len_i32(value: i32) -> usize {
encoded_len_u32(zigzag::encode_i32(value))
}
#[inline]
#[must_use]
pub const fn encoded_len_i64(value: i64) -> usize {
encoded_len_u64(zigzag::encode_i64(value))
}
#[inline]
pub fn encode_u16(value: u16, buf: &mut WriteBuf<'_>) -> Result<()> {
encode_raw(u64::from(value), buf)
}
#[inline]
pub fn encode_u32(value: u32, buf: &mut WriteBuf<'_>) -> Result<()> {
encode_raw(u64::from(value), buf)
}
#[inline]
pub fn encode_u64(value: u64, buf: &mut WriteBuf<'_>) -> Result<()> {
encode_raw(value, buf)
}
#[inline]
pub fn encode_usize(value: usize, buf: &mut WriteBuf<'_>) -> Result<()> {
encode_raw(value as u64, buf)
}
#[inline]
pub fn encode_i32(value: i32, buf: &mut WriteBuf<'_>) -> Result<()> {
encode_raw(u64::from(zigzag::encode_i32(value)), buf)
}
#[inline]
pub fn encode_i64(value: i64, buf: &mut WriteBuf<'_>) -> Result<()> {
encode_raw(zigzag::encode_i64(value), buf)
}
#[inline]
pub fn decode_u16(buf: &mut ReadBuf<'_>) -> Result<u16> {
decode_with(buf, MAX_LEN_U16, u16::BITS, false).map(|v| v as u16)
}
#[inline]
pub fn decode_u32(buf: &mut ReadBuf<'_>) -> Result<u32> {
decode_with(buf, MAX_LEN_U32, u32::BITS, false).map(|v| v as u32)
}
#[inline]
pub fn decode_u64(buf: &mut ReadBuf<'_>) -> Result<u64> {
decode_with(buf, MAX_LEN_U64, u64::BITS, false)
}
#[inline]
pub fn decode_usize(buf: &mut ReadBuf<'_>) -> Result<usize> {
decode_usize_with(buf, false)
}
#[inline]
pub fn decode_i32(buf: &mut ReadBuf<'_>) -> Result<i32> {
decode_u32(buf).map(zigzag::decode_i32)
}
#[inline]
pub fn decode_i64(buf: &mut ReadBuf<'_>) -> Result<i64> {
decode_u64(buf).map(zigzag::decode_i64)
}
#[inline]
pub fn decode_u16_canonical(buf: &mut ReadBuf<'_>) -> Result<u16> {
decode_with(buf, MAX_LEN_U16, u16::BITS, true).map(|v| v as u16)
}
#[inline]
pub fn decode_u32_canonical(buf: &mut ReadBuf<'_>) -> Result<u32> {
decode_with(buf, MAX_LEN_U32, u32::BITS, true).map(|v| v as u32)
}
#[inline]
pub fn decode_u64_canonical(buf: &mut ReadBuf<'_>) -> Result<u64> {
decode_with(buf, MAX_LEN_U64, u64::BITS, true)
}
#[inline]
pub fn decode_usize_canonical(buf: &mut ReadBuf<'_>) -> Result<usize> {
decode_usize_with(buf, true)
}
#[inline]
pub fn decode_i32_canonical(buf: &mut ReadBuf<'_>) -> Result<i32> {
decode_u32_canonical(buf).map(zigzag::decode_i32)
}
#[inline]
pub fn decode_i64_canonical(buf: &mut ReadBuf<'_>) -> Result<i64> {
decode_u64_canonical(buf).map(zigzag::decode_i64)
}
#[inline(always)]
fn encode_raw(mut value: u64, buf: &mut WriteBuf<'_>) -> Result<()> {
let dst = buf.spare_mut();
if dst.len() < MAX_LEN_U64 && dst.len() < encoded_len_u64(value) {
return Err(Error::BufferFull);
}
let mut i = 0;
while value >= 0x80 {
dst[i] = (value as u8) | 0x80;
value >>= 7;
i += 1;
}
dst[i] = value as u8;
buf.commit(i + 1);
Ok(())
}
#[inline(always)]
fn decode_with(buf: &mut ReadBuf<'_>, max_len: usize, bits: u32, canonical: bool) -> Result<u64> {
let (value, len) = decode_slice(buf.remaining_slice(), max_len, bits, canonical)?;
buf.consume(len);
Ok(value)
}
#[inline(always)]
fn decode_usize_with(buf: &mut ReadBuf<'_>, canonical: bool) -> Result<usize> {
let (value, len) = decode_slice(buf.remaining_slice(), MAX_LEN_U64, u64::BITS, canonical)?;
let value = usize::try_from(value).map_err(|_| Error::VarintOverflow)?;
buf.consume(len);
Ok(value)
}
#[inline(always)]
pub(crate) fn decode_slice(
bytes: &[u8],
max_len: usize,
bits: u32,
canonical: bool,
) -> Result<(u64, usize)> {
let mut value: u64 = 0;
let mut shift: u32 = 0;
for i in 0..max_len {
let Some(&byte) = bytes.get(i) else {
return Err(Error::UnexpectedEof);
};
value |= u64::from(byte & 0x7F) << shift;
if byte & 0x80 == 0 {
if i == max_len - 1 && u32::from(byte) >> (bits - shift) != 0 {
return Err(Error::VarintOverflow);
}
if canonical && byte == 0 && i != 0 {
return Err(Error::InvalidEncoding);
}
return Ok((value, i + 1));
}
shift += 7;
}
Err(Error::VarintOverflow)
}
#[cfg(test)]
mod tests {
use super::*;
fn round_trip_u32(value: u32) {
let mut storage = [0u8; MAX_LEN_U32];
let mut w = WriteBuf::new(&mut storage);
encode_u32(value, &mut w).unwrap();
let n = w.position();
assert_eq!(n, encoded_len_u32(value));
let mut r = ReadBuf::new(&storage[..n]);
assert_eq!(decode_u32(&mut r).unwrap(), value);
assert!(r.is_empty());
let mut r = ReadBuf::new(&storage[..n]);
assert_eq!(decode_u32_canonical(&mut r).unwrap(), value);
}
fn round_trip_u64(value: u64) {
let mut storage = [0u8; MAX_LEN_U64];
let mut w = WriteBuf::new(&mut storage);
encode_u64(value, &mut w).unwrap();
let n = w.position();
assert_eq!(n, encoded_len_u64(value));
let mut r = ReadBuf::new(&storage[..n]);
assert_eq!(decode_u64(&mut r).unwrap(), value);
assert!(r.is_empty());
let mut r = ReadBuf::new(&storage[..n]);
assert_eq!(decode_u64_canonical(&mut r).unwrap(), value);
}
#[test]
fn u32_boundaries() {
for &v in &[
0u32,
1,
127,
128,
16_383,
16_384,
2_097_151,
2_097_152,
u32::MAX,
] {
round_trip_u32(v);
}
}
#[test]
fn u64_boundaries() {
for &v in &[
0u64,
127,
128,
1 << 35,
1 << 56,
(1 << 63) - 1,
1 << 63,
u64::MAX,
] {
round_trip_u64(v);
}
}
#[test]
fn u16_round_trip_and_bounds() {
for v in [0u16, 1, 127, 128, 16_383, 16_384, u16::MAX] {
let mut storage = [0u8; MAX_LEN_U16];
let mut w = WriteBuf::new(&mut storage);
encode_u16(v, &mut w).unwrap();
let n = w.position();
assert_eq!(n, encoded_len_u16(v));
let mut r = ReadBuf::new(&storage[..n]);
assert_eq!(decode_u16(&mut r).unwrap(), v);
let mut r = ReadBuf::new(&storage[..n]);
assert_eq!(decode_u16_canonical(&mut r).unwrap(), v);
}
let mut r = ReadBuf::new(&[0xFF, 0xFF, 0x04]);
assert_eq!(decode_u16(&mut r), Err(Error::VarintOverflow));
let mut r = ReadBuf::new(&[0x80, 0x80, 0x80, 0x00]);
assert_eq!(decode_u16(&mut r), Err(Error::VarintOverflow));
}
#[test]
fn signed_round_trip() {
for v in [0i32, 1, -1, 63, -64, 64, -65, i32::MAX, i32::MIN] {
let mut storage = [0u8; MAX_LEN_U32];
let mut w = WriteBuf::new(&mut storage);
encode_i32(v, &mut w).unwrap();
let n = w.position();
assert_eq!(n, encoded_len_i32(v));
assert_eq!(decode_i32(&mut ReadBuf::new(&storage[..n])).unwrap(), v);
assert_eq!(
decode_i32_canonical(&mut ReadBuf::new(&storage[..n])).unwrap(),
v
);
}
for v in [0i64, 1, -1, i64::from(i32::MIN) - 1, i64::MAX, i64::MIN] {
let mut storage = [0u8; MAX_LEN_U64];
let mut w = WriteBuf::new(&mut storage);
encode_i64(v, &mut w).unwrap();
let n = w.position();
assert_eq!(n, encoded_len_i64(v));
assert_eq!(decode_i64(&mut ReadBuf::new(&storage[..n])).unwrap(), v);
assert_eq!(
decode_i64_canonical(&mut ReadBuf::new(&storage[..n])).unwrap(),
v
);
}
}
#[test]
fn usize_round_trip() {
for v in [0usize, 1, 300, usize::MAX] {
let mut storage = [0u8; MAX_LEN_U64];
let mut w = WriteBuf::new(&mut storage);
encode_usize(v, &mut w).unwrap();
let n = w.position();
assert_eq!(n, encoded_len_usize(v));
assert_eq!(decode_usize(&mut ReadBuf::new(&storage[..n])).unwrap(), v);
assert_eq!(
decode_usize_canonical(&mut ReadBuf::new(&storage[..n])).unwrap(),
v
);
}
}
#[test]
fn usize_too_wide_for_target_overflows() {
let bytes = [0x80, 0x80, 0x80, 0x80, 0x10];
let mut r = ReadBuf::new(&bytes);
let result = decode_usize(&mut r);
if usize::BITS > 32 {
assert_eq!(result.map(|v| v as u64), Ok(1u64 << 32));
} else {
assert_eq!(result, Err(Error::VarintOverflow));
assert_eq!(r.position(), 0);
}
}
#[test]
fn decode_u32_overflow_on_sixth_byte() {
let mut buf = ReadBuf::new(&[0x80, 0x80, 0x80, 0x80, 0x80, 0x01]);
assert_eq!(decode_u32(&mut buf), Err(Error::VarintOverflow));
assert_eq!(buf.position(), 0, "cursor is unchanged on error");
}
#[test]
fn decode_u32_overflow_in_final_byte() {
for last in [0x10u8, 0x1F, 0x7F] {
let bytes = [0xFF, 0xFF, 0xFF, 0xFF, last];
let mut buf = ReadBuf::new(&bytes);
assert_eq!(decode_u32(&mut buf), Err(Error::VarintOverflow));
}
let mut buf = ReadBuf::new(&[0xFF, 0xFF, 0xFF, 0xFF, 0x0F]);
assert_eq!(decode_u32(&mut buf), Ok(u32::MAX));
}
#[test]
fn decode_u64_final_byte_rules() {
let mut bytes = [0xFFu8; MAX_LEN_U64];
for last in [0x02u8, 0x03, 0x7F, 0x81, 0x80] {
bytes[9] = last;
let mut buf = ReadBuf::new(&bytes);
assert_eq!(
decode_u64(&mut buf),
Err(Error::VarintOverflow),
"last={last:#x}"
);
assert_eq!(buf.position(), 0);
}
bytes[9] = 0x01;
assert_eq!(decode_u64(&mut ReadBuf::new(&bytes)), Ok(u64::MAX));
bytes[9] = 0x00;
assert_eq!(decode_u64(&mut ReadBuf::new(&bytes)), Ok(u64::MAX >> 1));
assert_eq!(
decode_u64_canonical(&mut ReadBuf::new(&bytes)),
Err(Error::InvalidEncoding)
);
}
#[test]
fn decode_truncated_leaves_cursor_unchanged() {
let mut buf = ReadBuf::new(&[0x80, 0x80]);
assert_eq!(decode_u32(&mut buf), Err(Error::UnexpectedEof));
assert_eq!(buf.position(), 0);
let mut buf = ReadBuf::new(&[]);
assert_eq!(decode_u64(&mut buf), Err(Error::UnexpectedEof));
}
#[test]
fn streaming_retry_after_eof_succeeds() {
let wire = [0xAC, 0x02, 0x07];
let mut filled = 0;
let value = loop {
let mut r = ReadBuf::new(&wire[..filled]);
match decode_u32(&mut r) {
Ok(v) => {
assert_eq!(r.position(), 2);
break v;
}
Err(Error::UnexpectedEof) => {
assert_eq!(r.position(), 0);
filled += 1;
}
Err(e) => panic!("unexpected {e}"),
}
};
assert_eq!(value, 300);
assert_eq!(filled, 2);
}
#[test]
fn overlong_forms() {
let cases: [&[u8]; 4] = [
&[0x80, 0x00],
&[0x81, 0x00],
&[0x80, 0x80, 0x80, 0x80, 0x00],
&[0xFF, 0x80, 0x00],
];
for bytes in cases {
assert!(decode_u32(&mut ReadBuf::new(bytes)).is_ok());
let mut r = ReadBuf::new(bytes);
assert_eq!(decode_u32_canonical(&mut r), Err(Error::InvalidEncoding));
assert_eq!(r.position(), 0);
}
assert_eq!(decode_u32_canonical(&mut ReadBuf::new(&[0x00])), Ok(0));
assert_eq!(
decode_usize_canonical(&mut ReadBuf::new(&[0x80, 0x00])),
Err(Error::InvalidEncoding)
);
}
#[test]
fn encode_is_atomic_on_buffer_full() {
let mut storage = [0xEEu8; 1];
let mut w = WriteBuf::new(&mut storage);
assert_eq!(encode_u32(300, &mut w), Err(Error::BufferFull));
assert_eq!(w.position(), 0);
assert_eq!(storage, [0xEE], "nothing is written on error");
let mut storage = [0u8; 9];
let mut w = WriteBuf::new(&mut storage);
assert_eq!(encode_u64(u64::MAX, &mut w), Err(Error::BufferFull));
assert_eq!(w.position(), 0);
let mut storage = [0u8; 0];
let mut w = WriteBuf::new(&mut storage);
assert_eq!(encode_u16(0, &mut w), Err(Error::BufferFull));
assert_eq!(encode_i32(0, &mut w), Err(Error::BufferFull));
assert_eq!(encode_i64(0, &mut w), Err(Error::BufferFull));
assert_eq!(encode_usize(0, &mut w), Err(Error::BufferFull));
}
#[test]
fn exact_wire_bytes() {
let mut storage = [0u8; MAX_LEN_U64];
let mut w = WriteBuf::new(&mut storage);
encode_u64(0x80, &mut w).unwrap();
assert_eq!(w.written(), &[0x80, 0x01]);
let mut w = WriteBuf::new(&mut storage);
encode_u32(u32::MAX, &mut w).unwrap();
assert_eq!(w.written(), &[0xFF, 0xFF, 0xFF, 0xFF, 0x0F]);
}
}