Skip to main content

qubit_io/wrappers/
tee_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 mirrors accepted bytes into a branch writer.
18///
19/// `TeeWriter` writes to the primary writer first, then writes exactly the
20/// accepted bytes into the branch writer with [`Write::write_all`]. If the
21/// branch writer fails, the primary writer may already have accepted bytes and
22/// the branch error is returned.
23///
24/// # Examples
25/// ```
26/// use std::io::Write;
27///
28/// use qubit_io::TeeWriter;
29///
30/// let primary = Vec::new();
31/// let branch = Vec::new();
32/// let mut writer = TeeWriter::new(primary, branch);
33///
34/// writer.write_all(b"abc")?;
35/// let (primary, branch) = writer.into_inner();
36///
37/// assert_eq!(b"abc", primary.as_slice());
38/// assert_eq!(b"abc", branch.as_slice());
39/// # Ok::<(), std::io::Error>(())
40/// ```
41pub struct TeeWriter<P, B> {
42    primary: P,
43    branch: B,
44}
45
46impl<P, B> TeeWriter<P, B> {
47    /// Creates a tee writer.
48    ///
49    /// # Parameters
50    /// - `primary`: Primary destination writer.
51    /// - `branch`: Secondary writer that mirrors accepted bytes.
52    ///
53    /// # Returns
54    /// A new tee writer.
55    #[inline]
56    pub fn new(primary: P, branch: B) -> Self {
57        Self { primary, branch }
58    }
59
60    /// Returns an immutable reference to the primary writer.
61    ///
62    /// # Returns
63    /// The primary writer reference.
64    #[inline]
65    pub fn primary_ref(&self) -> &P {
66        &self.primary
67    }
68
69    /// Returns a mutable reference to the primary writer.
70    ///
71    /// # Returns
72    /// The primary writer reference.
73    #[inline]
74    pub fn primary_mut(&mut self) -> &mut P {
75        &mut self.primary
76    }
77
78    /// Returns an immutable reference to the branch writer.
79    ///
80    /// # Returns
81    /// The branch writer reference.
82    #[inline]
83    pub fn branch_ref(&self) -> &B {
84        &self.branch
85    }
86
87    /// Returns a mutable reference to the branch writer.
88    ///
89    /// # Returns
90    /// The branch writer reference.
91    #[inline]
92    pub fn branch_mut(&mut self) -> &mut B {
93        &mut self.branch
94    }
95
96    /// Consumes this wrapper and returns both wrapped writers.
97    ///
98    /// # Returns
99    /// A tuple containing the primary writer and branch writer.
100    #[inline]
101    pub fn into_inner(self) -> (P, B) {
102        (self.primary, self.branch)
103    }
104}
105
106impl<P, B> Write for TeeWriter<P, B>
107where
108    P: Write,
109    B: Write,
110{
111    fn write(&mut self, buffer: &[u8]) -> Result<usize> {
112        let count = self.primary.write(buffer)?;
113        self.branch.write_all(&buffer[..count])?;
114        Ok(count)
115    }
116
117    fn flush(&mut self) -> Result<()> {
118        self.primary.flush()?;
119        self.branch.flush()
120    }
121}
122
123impl<P, B> Seek for TeeWriter<P, B>
124where
125    P: Seek,
126    B: Seek,
127{
128    fn seek(&mut self, position: SeekFrom) -> Result<u64> {
129        let primary_position = self.primary.seek(position)?;
130        self.branch.seek(SeekFrom::Start(primary_position))?;
131        Ok(primary_position)
132    }
133}