pub const MAX_LENGTH: u32 = 0xFF_FF_FF;
#[inline]
pub fn try_read_length(input: &[u8]) -> Option<(u32, usize)> {
let first = *input.first()?;
match first >> 6 {
0 => Some((u32::from(first & 0x3F), 1)),
1 if input.len() > 1 => Some(((u32::from(first & 0x3F) << 8) | u32::from(input[1]), 2)),
2 if input.len() >= 5 => {
let len = u32::from_be_bytes(input[1..5].try_into().ok()?);
Some((len, 5))
}
_ => None,
}
}
#[inline]
pub fn try_write_length(length: u32, output: &mut [u8]) -> Option<usize> {
if length > MAX_LENGTH {
return None;
}
if length < 1 << 6 {
let byte = output.first_mut()?;
*byte = length as u8 & 0x3F;
return Some(1);
}
if length < 1 << 14 {
if output.len() < 2 {
return None;
}
output[0] = (((length >> 8) & 0x3F) as u8) | (1 << 6);
output[1] = length as u8;
return Some(2);
}
if output.len() < 5 {
return None;
}
output[0] = 2 << 6;
output[1..5].copy_from_slice(&length.to_be_bytes());
Some(5)
}
#[cfg(test)]
mod tests {
use super::{MAX_LENGTH, try_read_length, try_write_length};
#[test]
fn roundtrip_all_buckets() {
for len in [0u32, 63, 64, 0x3F_FF, 0x40_00, MAX_LENGTH] {
let mut buf = [0u8; 5];
let written = try_write_length(len, &mut buf).unwrap();
let (decoded, read) = try_read_length(&buf[..written]).unwrap();
assert_eq!((decoded, read), (len, written), "len={len}");
}
}
#[test]
fn encode_rejects_overflow_and_short_output() {
let mut buf = [0u8; 5];
assert_eq!(try_write_length(MAX_LENGTH + 1, &mut buf), None);
assert_eq!(try_write_length(64, &mut buf[..1]), None);
assert_eq!(try_write_length(16_384, &mut buf[..4]), None);
}
#[test]
fn decode_rejects_truncated_and_bad_prefix() {
assert_eq!(try_read_length(&[]), None);
assert_eq!(try_read_length(&[1 << 6]), None);
assert_eq!(try_read_length(&[2 << 6, 0, 0]), None);
assert_eq!(try_read_length(&[0xC0]), None);
}
}