qubit_io/wrappers/
limit_reader.rs

1// =============================================================================
2//    Copyright (c) 2026 Haixing Hu.
3//
4//    SPDX-License-Identifier: Apache-2.0
5//
6//    Licensed under the Apache License, Version 2.0.
7// =============================================================================
8use std::io::{
9    BufRead,
10    Read,
11    Result,
12};
13
14/// Reader wrapper that exposes at most a fixed number of bytes.
15///
16/// `LimitReader` is useful when a parser must consume a bounded section of a
17/// larger stream without relying on the caller to provide a pre-sliced buffer.
18/// Once the remaining limit reaches zero, reads return `Ok(0)` without touching
19/// the inner reader.
20///
21/// # Examples
22/// ```
23/// use std::io::{
24///     Cursor,
25///     Read,
26/// };
27///
28/// use qubit_io::LimitReader;
29///
30/// let mut reader = LimitReader::new(Cursor::new(b"abcdef"), 3);
31/// let mut data = Vec::new();
32/// reader.read_to_end(&mut data)?;
33///
34/// assert_eq!(b"abc", data.as_slice());
35/// assert_eq!(0, reader.remaining());
36/// # Ok::<(), std::io::Error>(())
37/// ```
38pub struct LimitReader<R> {
39    inner: R,
40    remaining: u64,
41}
42
43impl<R> LimitReader<R> {
44    /// Creates a reader that exposes at most `limit` bytes from `inner`.
45    ///
46    /// # Parameters
47    /// - `inner`: Reader to wrap.
48    /// - `limit`: Maximum number of bytes that may be read through this
49    ///   wrapper.
50    ///
51    /// # Returns
52    /// A new limited reader.
53    #[inline]
54    pub fn new(inner: R, limit: u64) -> Self {
55        Self {
56            inner,
57            remaining: limit,
58        }
59    }
60
61    /// Returns the number of bytes still available through this wrapper.
62    ///
63    /// # Returns
64    /// Remaining readable byte count before the wrapper reports EOF.
65    #[inline]
66    pub fn remaining(&self) -> u64 {
67        self.remaining
68    }
69
70    /// Returns an immutable reference to the wrapped reader.
71    ///
72    /// # Returns
73    /// The wrapped reader reference.
74    #[inline]
75    pub fn get_ref(&self) -> &R {
76        &self.inner
77    }
78
79    /// Returns a mutable reference to the wrapped reader.
80    ///
81    /// # Returns
82    /// The wrapped reader reference.
83    #[inline]
84    pub fn get_mut(&mut self) -> &mut R {
85        &mut self.inner
86    }
87
88    /// Consumes this wrapper and returns the wrapped reader.
89    ///
90    /// # Returns
91    /// The wrapped reader.
92    #[inline]
93    pub fn into_inner(self) -> R {
94        self.inner
95    }
96}
97
98impl<R> Read for LimitReader<R>
99where
100    R: Read,
101{
102    fn read(&mut self, buffer: &mut [u8]) -> Result<usize> {
103        if self.remaining == 0 || buffer.is_empty() {
104            return Ok(0);
105        }
106        let requested = self.remaining.min(buffer.len() as u64) as usize;
107        let count = self.inner.read(&mut buffer[..requested])?;
108        self.remaining -= count as u64;
109        Ok(count)
110    }
111}
112
113impl<R> BufRead for LimitReader<R>
114where
115    R: BufRead,
116{
117    fn fill_buf(&mut self) -> Result<&[u8]> {
118        if self.remaining == 0 {
119            return Ok(&[]);
120        }
121        let buffer = self.inner.fill_buf()?;
122        let limit = self.remaining.min(buffer.len() as u64) as usize;
123        Ok(&buffer[..limit])
124    }
125
126    fn consume(&mut self, amount: usize) {
127        let count = self.remaining.min(amount as u64);
128        self.remaining -= count;
129        self.inner.consume(count as usize);
130    }
131}
qubit_io/wrappers/limit_reader.rs

qubit_io/wrappers/
limit_reader.rs
