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 {}