use bytes::{Buf, BufMut, Bytes};
use crate::proto::ProtoError;
pub fn encode_vlq(value: u32, buf: &mut impl BufMut) {
if value < 0x80 {
buf.put_u8(value as u8);
return;
}
let mut groups: [u8; 5] = [0; 5];
let mut n = 0usize;
let mut v = value;
while v > 0 {
groups[n] = (v & 0x7F) as u8;
n += 1;
v >>= 7;
}
for i in (0..n).rev() {
let cont = if i > 0 { 0x80 } else { 0x00 };
buf.put_u8(groups[i] | cont);
}
}
pub fn decode_vlq(buf: &mut Bytes) -> Result<u32, ProtoError> {
let mut result = 0u64; let mut bytes_read = 0u32;
loop {
if buf.is_empty() {
return Err(ProtoError::BufferTooShort { need: 1, have: 0 });
}
let byte = buf.get_u8();
result = (result << 7) | ((byte & 0x7F) as u64);
if result > u32::MAX as u64 {
return Err(ProtoError::EncodingFailed("VLQ value exceeds u32".into()));
}
if byte & 0x80 == 0 {
return Ok(result as u32);
}
bytes_read += 1;
if bytes_read >= 5 {
return Err(ProtoError::EncodingFailed(
"VLQ overflow: more than 5 bytes".into(),
));
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use bytes::{Bytes, BytesMut};
#[test]
fn vlq_single_byte() {
let mut buf = BytesMut::new();
encode_vlq(0x2A, &mut buf);
assert_eq!(&buf[..], &[0x2A]);
}
#[test]
fn vlq_two_bytes() {
let mut buf = BytesMut::new();
encode_vlq(300, &mut buf);
assert_eq!(&buf[..], &[0x82, 0x2C]);
}
#[test]
fn vlq_decode_single() {
let mut b = Bytes::from_static(&[0x2A]);
assert_eq!(decode_vlq(&mut b).unwrap(), 0x2A);
}
#[test]
fn vlq_decode_two_bytes() {
let mut b = Bytes::from_static(&[0x82, 0x2C]);
assert_eq!(decode_vlq(&mut b).unwrap(), 300);
}
#[test]
fn vlq_roundtrip_large() {
for v in [
0u32,
1,
127,
128,
255,
300,
16383,
16384,
2097151,
2097152,
u32::MAX,
] {
let mut buf = BytesMut::new();
encode_vlq(v, &mut buf);
let mut b = buf.freeze();
assert_eq!(decode_vlq(&mut b).unwrap(), v, "roundtrip failed for {v}");
}
}
#[test]
fn vlq_decode_truncated_returns_err() {
let mut b = Bytes::from_static(&[0x82]); assert!(decode_vlq(&mut b).is_err());
}
#[test]
fn vlq_decode_overflow_returns_err() {
let mut b = Bytes::from_static(&[0x90, 0x80, 0x80, 0x80, 0x00]);
assert!(decode_vlq(&mut b).is_err());
}
}