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}