compio_io/write/
mod.rs

1#[cfg(feature = "allocator_api")]
2use std::alloc::Allocator;
3use std::io::Cursor;
4
5use compio_buf::{BufResult, IntoInner, IoBuf, IoVectoredBuf, buf_try, t_alloc};
6
7use crate::IoResult;
8
9mod buf;
10#[macro_use]
11mod ext;
12
13pub use buf::*;
14pub use ext::*;
15
16/// # AsyncWrite
17///
18/// Async write with a ownership of a buffer
19pub trait AsyncWrite {
20    /// Write some bytes from the buffer into this source and return a
21    /// [`BufResult`], consisting of the buffer and a [`usize`] indicating how
22    /// many bytes were written.
23    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T>;
24
25    /// Like `write`, except that it write bytes from a buffer implements
26    /// [`IoVectoredBuf`] into the source.
27    ///
28    /// The default implementation will try to write from the buffers in order
29    /// as if they're concatenated. It will stop whenever the writer returns
30    /// an error, `Ok(0)`, or a length less than the length of the buf passed
31    /// in, meaning it's possible that not all contents are written. If
32    /// guaranteed full write is desired, it is recommended to use
33    /// [`AsyncWriteExt::write_vectored_all`] instead.
34    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
35        loop_write_vectored!(buf, total: usize, n, iter, loop self.write(iter),
36            break if n == 0 || n < iter.buf_len() {
37                Some(Ok(total))
38            } else {
39                None
40            }
41        )
42    }
43
44    /// Attempts to flush the object, ensuring that any buffered data reach
45    /// their destination.
46    async fn flush(&mut self) -> IoResult<()>;
47
48    /// Initiates or attempts to shut down this writer, returning success when
49    /// the I/O connection has completely shut down.
50    async fn shutdown(&mut self) -> IoResult<()>;
51}
52
53impl<A: AsyncWrite + ?Sized> AsyncWrite for &mut A {
54    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
55        (**self).write(buf).await
56    }
57
58    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
59        (**self).write_vectored(buf).await
60    }
61
62    async fn flush(&mut self) -> IoResult<()> {
63        (**self).flush().await
64    }
65
66    async fn shutdown(&mut self) -> IoResult<()> {
67        (**self).shutdown().await
68    }
69}
70
71impl<W: AsyncWrite + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator> AsyncWrite
72    for t_alloc!(Box, W, A)
73{
74    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
75        (**self).write(buf).await
76    }
77
78    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
79        (**self).write_vectored(buf).await
80    }
81
82    async fn flush(&mut self) -> IoResult<()> {
83        (**self).flush().await
84    }
85
86    async fn shutdown(&mut self) -> IoResult<()> {
87        (**self).shutdown().await
88    }
89}
90
91/// Write is implemented for `Vec<u8>` by appending to the vector. The vector
92/// will grow as needed.
93impl AsyncWrite for Vec<u8> {
94    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
95        self.extend_from_slice(buf.as_slice());
96        BufResult(Ok(buf.buf_len()), buf)
97    }
98
99    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
100        let len = buf.iter_buf().map(|b| b.buf_len()).sum();
101        self.reserve(len - self.len());
102        for buf in buf.iter_buf() {
103            self.extend_from_slice(buf.as_slice());
104        }
105        BufResult(Ok(len), buf)
106    }
107
108    async fn flush(&mut self) -> IoResult<()> {
109        Ok(())
110    }
111
112    async fn shutdown(&mut self) -> IoResult<()> {
113        Ok(())
114    }
115}
116
117/// # AsyncWriteAt
118///
119/// Async write with a ownership of a buffer and a position
120pub trait AsyncWriteAt {
121    /// Like [`AsyncWrite::write`], except that it writes at a specified
122    /// position.
123    async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T>;
124
125    /// Like [`AsyncWrite::write_vectored`], except that it writes at a
126    /// specified position.
127    async fn write_vectored_at<T: IoVectoredBuf>(
128        &mut self,
129        buf: T,
130        pos: u64,
131    ) -> BufResult<usize, T> {
132        loop_write_vectored!(buf, total: u64, n, iter, loop self.write_at(iter, pos + total),
133            break if n == 0 || n < iter.buf_len() {
134                Some(Ok(total as usize))
135            } else {
136                None
137            }
138        )
139    }
140}
141
142impl<A: AsyncWriteAt + ?Sized> AsyncWriteAt for &mut A {
143    async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
144        (**self).write_at(buf, pos).await
145    }
146
147    async fn write_vectored_at<T: IoVectoredBuf>(
148        &mut self,
149        buf: T,
150        pos: u64,
151    ) -> BufResult<usize, T> {
152        (**self).write_vectored_at(buf, pos).await
153    }
154}
155
156impl<W: AsyncWriteAt + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator> AsyncWriteAt
157    for t_alloc!(Box, W, A)
158{
159    async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
160        (**self).write_at(buf, pos).await
161    }
162
163    async fn write_vectored_at<T: IoVectoredBuf>(
164        &mut self,
165        buf: T,
166        pos: u64,
167    ) -> BufResult<usize, T> {
168        (**self).write_vectored_at(buf, pos).await
169    }
170}
171
172macro_rules! impl_write_at {
173    ($($(const $len:ident =>)? $ty:ty),*) => {
174        $(
175            impl<$(const $len: usize)?> AsyncWriteAt for $ty {
176                async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
177                    let pos = (pos as usize).min(self.len());
178                    let slice = buf.as_slice();
179                    let n = slice.len().min(self.len() - pos);
180                    self[pos..pos + n].copy_from_slice(&slice[..n]);
181                    BufResult(Ok(n), buf)
182                }
183            }
184        )*
185    }
186}
187
188impl_write_at!([u8], const LEN => [u8; LEN]);
189
190/// This implementation aligns the behavior of files. If `pos` is larger than
191/// the vector length, the vectored will be extended, and the extended area will
192/// be filled with 0.
193impl<#[cfg(feature = "allocator_api")] A: Allocator> AsyncWriteAt for t_alloc!(Vec, u8, A) {
194    async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
195        let pos = pos as usize;
196        let slice = buf.as_slice();
197        if pos <= self.len() {
198            let n = slice.len().min(self.len() - pos);
199            if n < slice.len() {
200                self.reserve(slice.len() - n);
201                self[pos..pos + n].copy_from_slice(&slice[..n]);
202                self.extend_from_slice(&slice[n..]);
203            } else {
204                self[pos..pos + n].copy_from_slice(slice);
205            }
206            BufResult(Ok(n), buf)
207        } else {
208            self.reserve(pos - self.len() + slice.len());
209            self.resize(pos, 0);
210            self.extend_from_slice(slice);
211            BufResult(Ok(slice.len()), buf)
212        }
213    }
214}
215
216impl<A: AsyncWriteAt> AsyncWrite for Cursor<A> {
217    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
218        let pos = self.position();
219        let (n, buf) = buf_try!(self.get_mut().write_at(buf, pos).await);
220        self.set_position(pos + n as u64);
221        BufResult(Ok(n), buf)
222    }
223
224    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
225        let pos = self.position();
226        let (n, buf) = buf_try!(self.get_mut().write_vectored_at(buf, pos).await);
227        self.set_position(pos + n as u64);
228        BufResult(Ok(n), buf)
229    }
230
231    async fn flush(&mut self) -> IoResult<()> {
232        Ok(())
233    }
234
235    async fn shutdown(&mut self) -> IoResult<()> {
236        Ok(())
237    }
238}
compio_io/write/mod.rs

compio_io/write/
mod.rs
