qubit_io/stream/
zig_zag_reader.rs

1/*******************************************************************************
2 *
3 *    Copyright (c) 2026 Haixing Hu.
4 *
5 *    SPDX-License-Identifier: Apache-2.0
6 *
7 *    Licensed under the Apache License, Version 2.0.
8 *
9 ******************************************************************************/
10
11use core::marker::PhantomData;
12use std::io::{
13    Error,
14    ErrorKind,
15    Read,
16    Result,
17    Seek,
18    SeekFrom,
19};
20
21use crate::ReadExt;
22use crate::codec::{
23    DecodePolicy,
24    Leb128DecodeError,
25    NonStrict,
26    Strict,
27    ZigZagCodec,
28};
29
30/// Reader wrapper for ZigZag + unsigned LEB128 integers.
31///
32/// # Target-width integers
33///
34/// `isize` methods use the current Rust target's pointer width. Prefer
35/// fixed-width integer methods such as `read_i64` for persistent files and
36/// cross-platform protocols.
37pub struct ZigZagReader<R, P = NonStrict> {
38    inner: R,
39    buffer: [u8; 19],
40    marker: PhantomData<fn() -> P>,
41}
42
43impl<R, P> ZigZagReader<R, P>
44where
45    P: DecodePolicy,
46{
47    /// Creates a ZigZag reader.
48    #[must_use]
49    #[inline]
50    pub const fn new(inner: R) -> Self {
51        Self {
52            inner,
53            buffer: [0; 19],
54            marker: PhantomData,
55        }
56    }
57
58    /// Returns whether this reader rejects non-canonical LEB128 encodings.
59    #[must_use]
60    #[inline]
61    pub const fn is_strict(&self) -> bool {
62        P::STRICT
63    }
64
65    /// Returns a shared reference to the underlying reader.
66    #[must_use]
67    #[inline]
68    pub const fn get_ref(&self) -> &R {
69        &self.inner
70    }
71
72    /// Returns an exclusive reference to the underlying reader.
73    #[must_use]
74    #[inline]
75    pub fn get_mut(&mut self) -> &mut R {
76        &mut self.inner
77    }
78
79    /// Consumes this wrapper and returns the underlying reader.
80    #[must_use]
81    #[inline]
82    pub fn into_inner(self) -> R {
83        self.inner
84    }
85}
86
87macro_rules! impl_read_value {
88    ($policy:ty, $method:ident, $ty:ty, $doc:literal) => {
89        #[doc = $doc]
90        #[inline]
91        pub fn $method(&mut self) -> Result<$ty> {
92            type Codec = ZigZagCodec<$ty, $policy>;
93
94            self.read_leb128::<$ty, { Codec::REQUIRED_MIN_BUFFER_LEN }, _>(|bytes| unsafe {
95                Codec::read_unchecked(bytes, 0)
96            })
97        }
98    };
99}
100
101macro_rules! impl_for_policy {
102    ($policy:ty) => {
103        impl<R> ZigZagReader<R, $policy>
104        where
105            R: Read,
106        {
107            impl_read_value!($policy, read_i8, i8, "Reads a ZigZag `i8`.");
108            impl_read_value!($policy, read_i16, i16, "Reads a ZigZag `i16`.");
109            impl_read_value!($policy, read_i32, i32, "Reads a ZigZag `i32`.");
110            impl_read_value!($policy, read_i64, i64, "Reads a ZigZag `i64`.");
111            impl_read_value!($policy, read_i128, i128, "Reads a ZigZag `i128`.");
112            impl_read_value!($policy, read_isize, isize, "Reads a ZigZag `isize`.");
113        }
114    };
115}
116
117impl<R, P> ZigZagReader<R, P>
118where
119    R: Read,
120    P: DecodePolicy,
121{
122    #[inline]
123    fn read_leb128<T, const N: usize, F>(&mut self, decode: F) -> Result<T>
124    where
125        F: FnOnce(&[u8; 19]) -> std::result::Result<(T, usize), Leb128DecodeError>,
126    {
127        debug_assert!(N <= self.buffer.len(), "ZigZag read length exceeds internal buffer");
128        for index in 0..N {
129            // SAFETY: `index` is produced by `0..N`, where `N` is a
130            // codec-declared length that fits the fixed internal buffer.
131            unsafe {
132                self.inner.read_exact_unchecked(&mut self.buffer, index, 1)?;
133            }
134            if read_byte(&self.buffer, index) & 0x80 == 0 {
135                return decode(&self.buffer)
136                    .map(|(value, _)| value)
137                    .map_err(map_leb128_decode_error);
138            }
139        }
140        decode(&self.buffer)
141            .map(|(value, _)| value)
142            .map_err(map_leb128_decode_error)
143    }
144}
145
146impl_for_policy!(NonStrict);
147impl_for_policy!(Strict);
148
149impl<R, P> Read for ZigZagReader<R, P>
150where
151    R: Read,
152{
153    /// Reads bytes from the wrapped reader.
154    ///
155    /// # Parameters
156    ///
157    /// - `buffer`: Destination byte buffer.
158    ///
159    /// # Returns
160    ///
161    /// Returns the number of bytes read.
162    ///
163    /// # Errors
164    ///
165    /// Returns the I/O error reported by the wrapped reader.
166    #[inline]
167    fn read(&mut self, buffer: &mut [u8]) -> Result<usize> {
168        self.inner.read(buffer)
169    }
170}
171
172impl<R, P> Seek for ZigZagReader<R, P>
173where
174    R: Seek,
175{
176    /// Seeks the wrapped reader.
177    ///
178    /// # Parameters
179    ///
180    /// - `position`: Target seek position.
181    ///
182    /// # Returns
183    ///
184    /// Returns the new stream position.
185    ///
186    /// # Errors
187    ///
188    /// Returns the seek error reported by the wrapped reader.
189    #[inline]
190    fn seek(&mut self, position: SeekFrom) -> Result<u64> {
191        self.inner.seek(position)
192    }
193}
194
195#[inline]
196fn map_leb128_decode_error(error: Leb128DecodeError) -> Error {
197    Error::new(ErrorKind::InvalidData, error)
198}
199
200/// Reads one byte from the internal ZigZag buffer without an extra bounds check.
201#[inline(always)]
202fn read_byte(buffer: &[u8; 19], index: usize) -> u8 {
203    debug_assert!(index < buffer.len(), "ZigZag read index exceeds internal buffer");
204    // SAFETY: `read_leb128` only calls this with an index produced by
205    // `0..N`, where N is a codec-declared length that fits `buffer`.
206    unsafe { *buffer.as_ptr().add(index) }
207}
qubit_io/stream/zig_zag_reader.rs

qubit_io/stream/
zig_zag_reader.rs
