qubit_io/wrappers/
limit_writer.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    Result,
12    Write,
13};
14
15/// Writer wrapper that accepts at most a fixed number of bytes.
16///
17/// Once the remaining limit reaches zero, writes return `Ok(0)` without
18/// touching the inner writer. Callers using [`Write::write_all`] will therefore
19/// receive the standard write-zero error when trying to write beyond the limit.
20///
21/// # Examples
22/// ```
23/// use std::io::Write;
24///
25/// use qubit_io::LimitWriter;
26///
27/// let mut writer = LimitWriter::new(Vec::new(), 3);
28///
29/// assert_eq!(3, writer.write(b"abcdef")?);
30/// assert_eq!(0, writer.remaining());
31/// assert_eq!(0, writer.write(b"x")?);
32/// assert_eq!(b"abc", writer.get_ref().as_slice());
33/// # Ok::<(), std::io::Error>(())
34/// ```
35pub struct LimitWriter<W> {
36    inner: W,
37    remaining: u64,
38}
39
40impl<W> LimitWriter<W> {
41    /// Creates a writer that accepts at most `limit` bytes.
42    ///
43    /// # Parameters
44    /// - `inner`: Writer to wrap.
45    /// - `limit`: Maximum number of bytes that may be written through this
46    ///   wrapper.
47    ///
48    /// # Returns
49    /// A new limited writer.
50    #[inline]
51    pub fn new(inner: W, limit: u64) -> Self {
52        Self {
53            inner,
54            remaining: limit,
55        }
56    }
57
58    /// Returns the number of bytes still accepted by this wrapper.
59    ///
60    /// # Returns
61    /// Remaining writable byte count before the wrapper reports zero writes.
62    #[inline]
63    pub fn remaining(&self) -> u64 {
64        self.remaining
65    }
66
67    /// Returns an immutable reference to the wrapped writer.
68    ///
69    /// # Returns
70    /// The wrapped writer reference.
71    #[inline]
72    pub fn get_ref(&self) -> &W {
73        &self.inner
74    }
75
76    /// Returns a mutable reference to the wrapped writer.
77    ///
78    /// # Returns
79    /// The wrapped writer reference.
80    #[inline]
81    pub fn get_mut(&mut self) -> &mut W {
82        &mut self.inner
83    }
84
85    /// Consumes this wrapper and returns the wrapped writer.
86    ///
87    /// # Returns
88    /// The wrapped writer.
89    #[inline]
90    pub fn into_inner(self) -> W {
91        self.inner
92    }
93}
94
95impl<W> Write for LimitWriter<W>
96where
97    W: Write,
98{
99    fn write(&mut self, buffer: &[u8]) -> Result<usize> {
100        if self.remaining == 0 || buffer.is_empty() {
101            return Ok(0);
102        }
103        let requested = self.remaining.min(buffer.len() as u64) as usize;
104        let count = self.inner.write(&buffer[..requested])?;
105        self.remaining -= count as u64;
106        Ok(count)
107    }
108
109    #[inline]
110    fn flush(&mut self) -> Result<()> {
111        self.inner.flush()
112    }
113}
qubit_io/wrappers/limit_writer.rs

qubit_io/wrappers/
limit_writer.rs
