prefix_uvarint/
io.rs

1//! Extensions to `std::io` traits to support reading/writing prefix varints.
2use std::io::{BufRead, Error, ErrorKind, Read, Result, Write};
3
4use crate::raw::{encode_multibyte_writer_skip_2, max_value, tag_prefix};
5use crate::{DecodeError, PrefixVarInt, MAX_1BYTE_TAG, MAX_LEN};
6
7impl From<DecodeError> for Error {
8    fn from(value: DecodeError) -> Self {
9        let kind = match value {
10            DecodeError::UnexpectedEob => ErrorKind::UnexpectedEof,
11            DecodeError::Overflow => ErrorKind::InvalidData,
12        };
13        Error::new(kind, value)
14    }
15}
16
17/// Prefix varint code a value and write it to `w`. Returns the number of bytes
18/// written.
19#[inline]
20pub fn write_prefix_varint<PV: PrefixVarInt>(v: PV, w: &mut impl Write) -> Result<usize> {
21    let v = v.to_prefix_varint_raw();
22    if v <= max_value(1) {
23        return w.write_all(&[v as u8]).map(|_| 1);
24    } else if v <= max_value(2) {
25        let tag_prefix = tag_prefix(2) >> 48;
26        let tagged = (v | tag_prefix) as u16;
27        return w.write_all(&tagged.to_be_bytes()).map(|_| 2);
28    }
29    encode_multibyte_writer_skip_2(v, w)
30}
31
32/// Read and decode a prefix varint value from `r`.
33/// Prefer `read_prefix_varint_buf()` wherever possible as it should be more efficient.
34#[inline]
35pub fn read_prefix_varint<PV: PrefixVarInt>(r: &mut impl Read) -> Result<PV> {
36    let mut buf = [0u8; MAX_LEN];
37    r.read_exact(&mut buf[..1])?;
38    let tag = buf[0];
39    if tag <= MAX_1BYTE_TAG {
40        return PV::from_prefix_varint_raw(tag.into()).ok_or_else(|| DecodeError::Overflow.into());
41    }
42    let rem = tag.leading_ones() as usize;
43    r.read_exact(&mut buf[1..(1 + rem)])?;
44    PV::decode_prefix_varint(buf.as_slice())
45        .map(|(v, _)| v)
46        .map_err(|e| e.into())
47}
48
49/// Read and decode a prefix varint value from `r`.
50#[inline]
51pub fn read_prefix_varint_buf<PV: PrefixVarInt>(r: &mut impl BufRead) -> Result<PV> {
52    let buf = r.fill_buf()?;
53    if buf.len() >= MAX_LEN {
54        let (v, len) = PV::decode_prefix_varint(buf).map_err(Error::from)?;
55        r.consume(len);
56        Ok(v)
57    } else {
58        read_prefix_varint(r)
59    }
60}