use std::io::{Read, Write};
use crate::error::{EvalErr, Result};
pub fn write_varint<W: Write>(w: &mut W, value: i64) -> std::io::Result<()> {
for leading_ones in 0..8 {
let total_value_bits = 7 + 7 * leading_ones;
let min_value = -(1i64 << (total_value_bits - 1));
let max_value = (1i64 << (total_value_bits - 1)) - 1;
if value < min_value || value > max_value {
continue; }
let unsigned_value = if value < 0 {
(value + (1i64 << total_value_bits)) as u64
} else {
value as u64
};
let first_byte = if leading_ones > 0 {
((1u8 << leading_ones) - 1) << (8 - leading_ones)
} else {
0
};
let high_bits = (unsigned_value >> (leading_ones * 8)) as u8;
let first_byte = first_byte | high_bits;
w.write_all(&[first_byte])?;
for i in (0..leading_ones).rev() {
let byte_val = (unsigned_value >> (i * 8)) as u8;
w.write_all(&[byte_val])?;
}
return Ok(());
}
panic!("Value too large to encode: {}", value);
}
fn varint_size(value: i64) -> usize {
for leading_ones in 0..8 {
let total_value_bits = 7 + 7 * leading_ones;
let min_value = -(1i64 << (total_value_bits - 1));
let max_value = (1i64 << (total_value_bits - 1)) - 1;
if value >= min_value && value <= max_value {
return leading_ones + 1;
}
}
panic!("Value too large to encode: {}", value);
}
pub fn read_varint<R: Read>(r: &mut R, strict: bool) -> Result<i64> {
let mut first_byte_buf = [0u8; 1];
r.read_exact(&mut first_byte_buf)
.map_err(|_| EvalErr::SerializationError)?;
let first_byte = first_byte_buf[0];
let leading_ones = (!first_byte).leading_zeros() as usize;
if leading_ones >= 8 {
return Err(EvalErr::SerializationError);
}
let bits_in_first_byte = 7 - leading_ones;
let total_value_bits = 7 + 7 * leading_ones;
let value_mask = (1u8 << bits_in_first_byte) - 1;
let mut unsigned_value = (first_byte & value_mask) as u64;
if leading_ones > 0 {
let mut extra_bytes = [0u8; 7];
r.read_exact(&mut extra_bytes[..leading_ones])
.map_err(|_| EvalErr::SerializationError)?;
for &byte in &extra_bytes[..leading_ones] {
unsigned_value = (unsigned_value << 8) | (byte as u64);
}
}
let sign_bit = 1u64 << (total_value_bits - 1);
let value = if unsigned_value >= sign_bit {
unsigned_value as i64 - (1i64 << total_value_bits)
} else {
unsigned_value as i64
};
if strict && varint_size(value) != leading_ones + 1 {
return Err(EvalErr::SerializationError);
}
Ok(value)
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
use std::io::Cursor;
fn encode_varint(value: i64) -> Vec<u8> {
let mut buf = Vec::new();
write_varint(&mut buf, value).unwrap();
buf
}
#[rstest]
#[case(0, vec![0x00])]
#[case(1, vec![0x01])]
#[case(-1, vec![0x7f])]
#[case(63, vec![0x3f])]
#[case(-64, vec![0x40])]
#[case(64, vec![0x80, 0x40])]
#[case(8191, vec![0x9f, 0xff])]
#[case(-65, vec![0xbf, 0xbf])]
#[case(-8192, vec![0xa0, 0x00])]
fn test_encode_varint(#[case] value: i64, #[case] expected: Vec<u8>) {
assert_eq!(encode_varint(value), expected);
}
#[rstest]
#[case(&[0x00], 0)]
#[case(&[0x01], 1)]
#[case(&[0x7f], -1)]
#[case(&[0x80, 0x40], 64)]
#[case(&[0x9f, 0xff], 8191)]
#[case(&[0xbf, 0xbf], -65)]
fn test_read_varint(#[case] bytes: &[u8], #[case] expected: i64) {
assert_eq!(
read_varint(&mut Cursor::new(bytes), false).unwrap(),
expected
);
}
#[test]
#[should_panic(expected = "Value too large to encode")]
fn test_encode_varint_rejects_too_large_positive() {
let _ = encode_varint(1i64 << 55); }
#[test]
#[should_panic(expected = "Value too large to encode")]
fn test_encode_varint_rejects_too_large_negative() {
let _ = encode_varint(-(1i64 << 55) - 1); }
#[test]
fn test_encode_varint_accepts_boundary_values() {
let max_value = (1i64 << 55) - 1; let min_value = -(1i64 << 55);
let _ = encode_varint(max_value);
let _ = encode_varint(min_value);
}
#[test]
fn test_decode_rejects_invalid_prefix() {
assert!(read_varint(&mut Cursor::new(&[0xff][..]), false).is_err());
assert!(read_varint(&mut Cursor::new(&[0xff, 0x00][..]), false).is_err());
}
#[test]
fn test_decode_rejects_truncated_multibyte() {
assert!(read_varint(&mut Cursor::new(&[0x80][..]), false).is_err());
assert!(read_varint(&mut Cursor::new(&[0xc0, 0x00][..]), false).is_err());
assert!(read_varint(&mut Cursor::new(&[0xe0, 0x00, 0x00][..]), false).is_err());
}
#[test]
fn test_decode_empty_input() {
assert!(read_varint(&mut Cursor::new(&[][..]), false).is_err());
}
#[test]
fn test_strict_decode_rejects_overlong_encodings() {
assert_eq!(
read_varint(&mut Cursor::new(&[0x80, 0x00][..]), false).unwrap(),
0
);
assert!(read_varint(&mut Cursor::new(&[0x80, 0x00][..]), true).is_err());
assert_eq!(
read_varint(&mut Cursor::new(&[0xbf, 0xff][..]), false).unwrap(),
-1
);
assert!(read_varint(&mut Cursor::new(&[0xbf, 0xff][..]), true).is_err());
}
#[test]
fn test_roundtrip_boundary_values() {
for val in [
0, 1, -1, 63, -64, 64, -65, 8191, -8192, 8192, -8193, 1048575, -1048576,
] {
let encoded = encode_varint(val);
let decoded = read_varint(&mut Cursor::new(&encoded), false).unwrap();
assert_eq!(val, decoded, "roundtrip failed for {val}");
}
}
}