Skip to main content

qubit_io/adapters/
tokio_output.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::pin::Pin;
10use std::task::Context;
11use std::task::Poll;
12
13use tokio::io::AsyncWrite;
14
15use crate::AsyncClose;
16use crate::AsyncOutput;
17use crate::traits::normalize_async_error;
18use crate::util::UncheckedSlice;
19
20/// Adapts a Tokio [`AsyncWrite`] value to Qubit's [`AsyncOutput`].
21///
22/// # Type Parameters
23///
24/// - `T`: Tokio writer type.
25#[must_use]
26#[repr(transparent)]
27pub struct TokioOutput<T> {
28    /// Tokio writer adapted as a Qubit output.
29    inner: T,
30}
31
32impl<T> TokioOutput<T> {
33    /// Creates an adapter around a Tokio writer.
34    ///
35    /// # Parameters
36    ///
37    /// - `inner`: Tokio writer to adapt.
38    ///
39    /// # Returns
40    ///
41    /// Returns a Qubit output adapter that owns `inner`.
42    #[inline(always)]
43    pub const fn new(inner: T) -> Self {
44        Self { inner }
45    }
46
47    /// Returns a shared reference to the wrapped writer.
48    ///
49    /// # Returns
50    ///
51    /// Returns the wrapped Tokio writer.
52    #[inline(always)]
53    #[must_use]
54    pub const fn get_ref(&self) -> &T {
55        &self.inner
56    }
57
58    /// Returns a mutable reference to the wrapped writer.
59    ///
60    /// # Returns
61    ///
62    /// Returns the wrapped Tokio writer with mutable access.
63    #[inline(always)]
64    #[must_use]
65    pub const fn get_mut(&mut self) -> &mut T {
66        &mut self.inner
67    }
68
69    /// Projects a pinned adapter to its pinned wrapped writer.
70    ///
71    /// # Returns
72    ///
73    /// Returns a pinned mutable reference to the wrapped writer without moving
74    /// it.
75    #[inline(always)]
76    #[must_use]
77    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> {
78        // SAFETY: The projection does not move `inner`, and the transparent
79        // adapter never exposes a way to replace a pinned inner value.
80        unsafe { self.map_unchecked_mut(|this| &mut this.inner) }
81    }
82
83    /// Consumes the adapter and returns the wrapped writer.
84    ///
85    /// # Returns
86    ///
87    /// Returns the owned Tokio writer.
88    #[inline(always)]
89    #[must_use]
90    pub fn into_inner(self) -> T {
91        self.inner
92    }
93}
94
95impl<T> AsyncOutput for TokioOutput<T>
96where
97    T: AsyncWrite,
98{
99    /// Byte item accepted by a Tokio writer.
100    type Item = u8;
101
102    /// Polls an indexed write through the wrapped Tokio writer.
103    ///
104    /// A zero-length request completes immediately without polling `inner`.
105    ///
106    /// # Parameters
107    ///
108    /// - `cx`: Task context used to register a wake-up.
109    /// - `input`: Source byte slice.
110    /// - `index`: Starting source index.
111    /// - `count`: Maximum number of bytes to write.
112    ///
113    /// # Returns
114    ///
115    /// Returns [`Poll::Pending`] when the writer is not ready. A ready result
116    /// contains the number of bytes accepted.
117    ///
118    /// # Errors
119    ///
120    /// Returns an I/O error reported by the wrapped writer. Invalid
121    /// asynchronous error kinds are normalized to
122    /// [`std::io::ErrorKind::InvalidData`].
123    ///
124    /// # Panics
125    ///
126    /// Panics in debug builds if the requested input range does not fit.
127    ///
128    /// # Safety
129    ///
130    /// The range `index..index + count` must be valid for `input`.
131    #[inline]
132    unsafe fn poll_write_unchecked(
133        self: Pin<&mut Self>,
134        cx: &mut Context<'_>,
135        input: &[u8],
136        index: usize,
137        count: usize,
138    ) -> Poll<std::io::Result<usize>> {
139        if count == 0 {
140            return Poll::Ready(Ok(0));
141        }
142        // SAFETY: The caller guarantees that the source range is valid.
143        let source = unsafe { UncheckedSlice::subslice(input, index, count) };
144        AsyncWrite::poll_write(self.get_pin_mut(), cx, source).map(|result| result.map_err(normalize_async_error))
145    }
146
147    /// Polls flushing through the wrapped Tokio writer.
148    ///
149    /// # Parameters
150    ///
151    /// - `cx`: Task context used to register a wake-up.
152    ///
153    /// # Returns
154    ///
155    /// Returns [`Poll::Pending`] while flushing is incomplete, otherwise a
156    /// ready success result.
157    ///
158    /// # Errors
159    ///
160    /// Returns an I/O error reported by the wrapped writer. Invalid
161    /// asynchronous error kinds are normalized to
162    /// [`std::io::ErrorKind::InvalidData`].
163    #[inline(always)]
164    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
165        AsyncWrite::poll_flush(self.get_pin_mut(), cx).map(|result| result.map_err(normalize_async_error))
166    }
167}
168
169impl<T> AsyncClose for TokioOutput<T>
170where
171    T: AsyncWrite,
172{
173    /// Polls closing through the wrapped Tokio writer.
174    ///
175    /// # Parameters
176    ///
177    /// - `cx`: Task context used to register a wake-up.
178    ///
179    /// # Returns
180    ///
181    /// Returns [`Poll::Pending`] while closing is incomplete, otherwise a
182    /// ready success result.
183    ///
184    /// # Errors
185    ///
186    /// Returns an I/O error reported by the wrapped writer. Invalid
187    /// asynchronous error kinds are normalized to
188    /// [`std::io::ErrorKind::InvalidData`].
189    #[inline(always)]
190    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
191        AsyncWrite::poll_shutdown(self.get_pin_mut(), cx).map(|result| result.map_err(normalize_async_error))
192    }
193}