Skip to main content

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