Skip to main content

qubit_io/ext/
output_ext.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// =============================================================================
8
9use std::io::{
10    Error,
11    ErrorKind,
12    Result,
13};
14
15use crate::Output;
16use crate::util::UncheckedSlice;
17
18/// Extension methods for [`Output`] values.
19///
20/// `OutputExt` keeps complete-write helpers outside the minimal [`Output`]
21/// trait. The methods are implemented for every item-oriented output,
22/// including `dyn Output` trait objects.
23pub trait OutputExt: Output {
24    /// Writes all items from the full input slice.
25    ///
26    /// This method repeatedly calls [`Output::write_unchecked`] until the full
27    /// slice is accepted. Interrupted writes are retried.
28    ///
29    /// # Parameters
30    /// - `input`: Source items.
31    ///
32    /// # Errors
33    /// Returns the output error reported by the implementation. Returns
34    /// [`ErrorKind::WriteZero`] if the output accepts zero items before the
35    /// slice is complete. Returns [`ErrorKind::InvalidData`] if the output
36    /// reports accepting more items than requested.
37    #[inline(always)]
38    fn write_all(&mut self, input: &[Self::Item]) -> Result<()> {
39        // SAFETY: The full input slice is a valid source range.
40        unsafe { self.write_all_unchecked(input, 0, input.len()) }
41    }
42
43    /// Writes all items from an indexed input range without checking the range.
44    ///
45    /// This method repeatedly calls [`Output::write_unchecked`] until all
46    /// `count` items are accepted. Interrupted writes are retried. A zero
47    /// progress report before the range is complete is converted to
48    /// [`ErrorKind::WriteZero`].
49    ///
50    /// # Parameters
51    /// - `input`: Source storage.
52    /// - `index`: Start index inside `input`.
53    /// - `count`: Number of items to write.
54    ///
55    /// # Errors
56    /// Returns the output error reported by the implementation. Returns
57    /// [`ErrorKind::WriteZero`] if the implementation accepts zero items before
58    /// the requested range is complete. Returns [`ErrorKind::InvalidData`] if
59    /// the implementation reports accepting more items than requested.
60    ///
61    /// # Safety
62    /// The caller must guarantee that `index..index + count` is a valid range
63    /// inside `input` and that the addition does not overflow.
64    unsafe fn write_all_unchecked(
65        &mut self,
66        input: &[Self::Item],
67        index: usize,
68        count: usize,
69    ) -> Result<()> {
70        debug_assert!(
71            UncheckedSlice::range_fits(input.len(), index, count),
72            "unchecked write-all range exceeds input buffer"
73        );
74        let mut written = 0;
75        while written < count {
76            let remaining = count - written;
77            // SAFETY: The caller guarantees the original source range is
78            // valid; `written < count`, so this suffix remains inside it.
79            match unsafe {
80                self.write_unchecked(input, index + written, remaining)
81            } {
82                Ok(0) => {
83                    return Err(Error::new(
84                        ErrorKind::WriteZero,
85                        "failed to write whole output range",
86                    ));
87                }
88                Ok(progress) => {
89                    if progress > remaining {
90                        return Err(Error::new(
91                            ErrorKind::InvalidData,
92                            format!(
93                                "writer reported {progress} items for a {remaining}-item range"
94                            ),
95                        ));
96                    }
97                    written += progress;
98                }
99                Err(error) if error.kind() == ErrorKind::Interrupted => {}
100                Err(error) => return Err(error),
101            }
102        }
103        Ok(())
104    }
105}
106
107impl<T> OutputExt for T where T: Output + ?Sized {}