Skip to main content

qubit_io/async_io/
write_fully_future.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::future::Future;
10use std::io::Error;
11use std::io::ErrorKind;
12use std::io::Result;
13use std::pin::Pin;
14use std::task::Context;
15use std::task::Poll;
16
17use super::read_fully_future::MAX_READY_OPERATIONS_PER_POLL;
18use crate::AsyncOutput;
19
20/// Future that writes every item from its source slice.
21///
22/// Items accepted before cancellation remain written. The accepted count is
23/// observable through [`Self::items_written`].
24///
25/// To preserve executor fairness, one outer poll performs a bounded number of
26/// successful inner writes. The future self-wakes and returns [`Poll::Pending`]
27/// when that budget is exhausted before completion.
28///
29/// # Panics
30///
31/// [`Future::poll`] panics when called again after this future has returned
32/// [`Poll::Ready`].
33///
34/// # Type Parameters
35///
36/// - `'a`: Shared lifetime of the output borrow and source slice.
37/// - `O`: Asynchronous output type.
38#[must_use = "futures do nothing unless polled"]
39pub struct WriteFullyFuture<'a, O>
40where
41    O: AsyncOutput + ?Sized,
42{
43    /// Output being written.
44    output: Pin<&'a mut O>,
45    /// Source whose items must all be written.
46    input: &'a [O::Item],
47    /// Number of items written so far.
48    written: usize,
49    /// Whether the write operation has completed.
50    completed: bool,
51}
52
53impl<'a, O> WriteFullyFuture<'a, O>
54where
55    O: AsyncOutput + ?Sized,
56{
57    /// Creates a write-fully future from a pinned output.
58    ///
59    /// # Parameters
60    ///
61    /// - `output`: Pinned asynchronous output.
62    /// - `input`: Source storage.
63    ///
64    /// # Returns
65    ///
66    /// Returns a future that resolves after every item has been accepted.
67    #[inline(always)]
68    pub const fn new(output: Pin<&'a mut O>, input: &'a [O::Item]) -> Self {
69        Self {
70            output,
71            input,
72            written: 0,
73            completed: false,
74        }
75    }
76
77    /// Returns the number of items written so far.
78    ///
79    /// # Returns
80    ///
81    /// Returns the progress retained across polls.
82    #[inline(always)]
83    #[must_use]
84    pub const fn items_written(&self) -> usize {
85        self.written
86    }
87}
88
89impl<O> Future for WriteFullyFuture<'_, O>
90where
91    O: AsyncOutput + ?Sized,
92{
93    /// Result produced when the write-fully operation becomes ready.
94    type Output = Result<()>;
95
96    /// Polls the write-fully operation.
97    ///
98    /// # Parameters
99    ///
100    /// - `cx`: Task context used to register a wake-up.
101    ///
102    /// # Returns
103    ///
104    /// Returns [`Poll::Pending`] while more output capacity is needed, or
105    /// [`Poll::Ready`] after all items are written or an error occurs.
106    ///
107    /// # Errors
108    ///
109    /// Returns [`ErrorKind::WriteZero`] if the output accepts no item before
110    /// completion. Other errors are propagated from the output.
111    ///
112    /// # Panics
113    ///
114    /// Panics when polled after returning [`Poll::Ready`].
115    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
116        let this = self.get_mut();
117        assert!(!this.completed, "WriteFullyFuture polled after completion");
118        let mut ready_operations = 0_usize;
119        while this.written < this.input.len() {
120            let remaining = &this.input[this.written..];
121            match this.output.as_mut().poll_write(cx, remaining) {
122                Poll::Ready(Ok(0)) => {
123                    this.completed = true;
124                    return Poll::Ready(Err(Error::new(
125                        ErrorKind::WriteZero,
126                        "failed to write whole output range",
127                    )));
128                }
129                Poll::Ready(Ok(written)) => {
130                    this.written += written;
131                    ready_operations += 1;
132                    if this.written < this.input.len() && ready_operations >= MAX_READY_OPERATIONS_PER_POLL {
133                        cx.waker().wake_by_ref();
134                        return Poll::Pending;
135                    }
136                }
137                Poll::Ready(Err(error)) => {
138                    this.completed = true;
139                    return Poll::Ready(Err(error));
140                }
141                Poll::Pending => return Poll::Pending,
142            }
143        }
144        this.completed = true;
145        Poll::Ready(Ok(()))
146    }
147}