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

qubit_io/wrappers/
limit_writer.rs
