Skip to main content

qubit_io/wrappers/
counting_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    Seek,
13    SeekFrom,
14    Write,
15};
16
17/// Writer wrapper that counts successfully written bytes.
18///
19/// The count is increased only after the wrapped writer reports a successful
20/// byte count. Failed writes do not change the counter.
21///
22/// # Examples
23/// ```
24/// use std::io::Write;
25///
26/// use qubit_io::CountingWriter;
27///
28/// let mut writer = CountingWriter::new(Vec::new());
29/// writer.write_all(b"abc")?;
30///
31/// assert_eq!(3, writer.bytes_written());
32/// assert_eq!(b"abc", writer.get_ref().as_slice());
33/// # Ok::<(), std::io::Error>(())
34/// ```
35pub struct CountingWriter<W> {
36    inner: W,
37    bytes_written: u64,
38}
39
40impl<W> CountingWriter<W> {
41    /// Creates a counting writer.
42    ///
43    /// # Parameters
44    /// - `inner`: Writer to wrap.
45    ///
46    /// # Returns
47    /// A new counting writer with a zero byte count.
48    #[inline]
49    pub fn new(inner: W) -> Self {
50        Self {
51            inner,
52            bytes_written: 0,
53        }
54    }
55
56    /// Returns the number of bytes successfully written through this wrapper.
57    ///
58    /// # Returns
59    /// Total byte count. The value saturates at [`u64::MAX`].
60    #[inline]
61    pub fn bytes_written(&self) -> u64 {
62        self.bytes_written
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 CountingWriter<W>
94where
95    W: Write,
96{
97    fn write(&mut self, buffer: &[u8]) -> Result<usize> {
98        let count = self.inner.write(buffer)?;
99        self.bytes_written = self.bytes_written.saturating_add(count as u64);
100        Ok(count)
101    }
102
103    #[inline]
104    fn flush(&mut self) -> Result<()> {
105        self.inner.flush()
106    }
107}
108
109impl<W> Seek for CountingWriter<W>
110where
111    W: Seek,
112{
113    #[inline]
114    fn seek(&mut self, position: SeekFrom) -> Result<u64> {
115        self.inner.seek(position)
116    }
117}