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

qubit_io/wrappers/
limit_reader.rs
