vortex_io/
write.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::future::Future;
5use std::future::ready;
6use std::io::Cursor;
7use std::io::Write;
8use std::io::{self};
9
10use futures::AsyncWrite;
11use futures::AsyncWriteExt;
12use vortex_buffer::ByteBufferMut;
13
14use crate::IoBuf;
15
16/// The unified write trait for Vortex I/O sinks.
17///
18/// Implementations and their returned futures are `Send`, mirroring [`VortexReadAt`], so writers
19/// can be driven from multi-threaded executors and moved into spawned tasks.
20///
21/// [`VortexReadAt`]: crate::VortexReadAt
22pub trait VortexWrite: Send {
23    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> + Send;
24    fn flush(&mut self) -> impl Future<Output = io::Result<()>> + Send;
25    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> + Send;
26}
27
28impl VortexWrite for Vec<u8> {
29    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> + Send {
30        self.extend_from_slice(buffer.as_slice());
31        ready(Ok(buffer))
32    }
33
34    fn flush(&mut self) -> impl Future<Output = io::Result<()>> + Send {
35        ready(Ok(()))
36    }
37
38    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> + Send {
39        ready(Ok(()))
40    }
41}
42
43impl VortexWrite for ByteBufferMut {
44    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> + Send {
45        self.extend_from_slice(buffer.as_slice());
46        ready(Ok(buffer))
47    }
48
49    fn flush(&mut self) -> impl Future<Output = io::Result<()>> + Send {
50        ready(Ok(()))
51    }
52
53    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> + Send {
54        ready(Ok(()))
55    }
56}
57
58impl<T: Send> VortexWrite for Cursor<T>
59where
60    Cursor<T>: Write,
61{
62    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> + Send {
63        ready(Write::write_all(self, buffer.as_slice()).map(|_| buffer))
64    }
65
66    fn flush(&mut self) -> impl Future<Output = io::Result<()>> + Send {
67        ready(Write::flush(self))
68    }
69
70    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> + Send {
71        ready(Write::flush(self))
72    }
73}
74
75impl<W: VortexWrite> VortexWrite for futures::io::Cursor<W> {
76    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> + Send {
77        self.set_position(self.position() + buffer.as_slice().len() as u64);
78        VortexWrite::write_all(self.get_mut(), buffer)
79    }
80
81    fn flush(&mut self) -> impl Future<Output = io::Result<()>> + Send {
82        VortexWrite::flush(self.get_mut())
83    }
84
85    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> + Send {
86        VortexWrite::shutdown(self.get_mut())
87    }
88}
89
90impl<W: VortexWrite> VortexWrite for &mut W {
91    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> + Send {
92        (*self).write_all(buffer)
93    }
94
95    fn flush(&mut self) -> impl Future<Output = io::Result<()>> + Send {
96        (*self).flush()
97    }
98
99    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> + Send {
100        (*self).shutdown()
101    }
102}
103
104impl VortexWrite for async_fs::File {
105    async fn write_all<B: IoBuf>(&mut self, buffer: B) -> io::Result<B> {
106        AsyncWriteExt::write_all(self, buffer.as_slice()).await?;
107        Ok(buffer)
108    }
109
110    fn flush(&mut self) -> impl Future<Output = io::Result<()>> + Send {
111        AsyncWriteExt::flush(self)
112    }
113
114    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> + Send {
115        AsyncWriteExt::close(self)
116    }
117}
118
119/// An adapter to use an `AsyncWrite` as a `VortexWrite`.
120pub struct AsyncWriteAdapter<W: AsyncWrite>(pub W);
121impl<W: AsyncWrite + Unpin + Send> VortexWrite for AsyncWriteAdapter<W> {
122    async fn write_all<B: IoBuf>(&mut self, buffer: B) -> io::Result<B> {
123        self.0.write_all(buffer.as_slice()).await?;
124        Ok(buffer)
125    }
126
127    async fn flush(&mut self) -> io::Result<()> {
128        self.0.flush().await
129    }
130
131    async fn shutdown(&mut self) -> io::Result<()> {
132        self.0.close().await
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use rstest::rstest;
139    use vortex_buffer::ByteBufferMut;
140
141    use super::*;
142
143    /// A generic writer's futures must stay `Send` so writers can be driven on multi-threaded
144    /// executors; this only compiles while that holds.
145    #[tokio::test]
146    async fn test_write_futures_are_send() -> io::Result<()> {
147        async fn write_generic<W: VortexWrite>(mut writer: W) -> io::Result<W> {
148            VortexWrite::write_all(&mut writer, vec![1, 2, 3]).await?;
149            VortexWrite::flush(&mut writer).await?;
150            VortexWrite::shutdown(&mut writer).await?;
151            Ok(writer)
152        }
153
154        fn assert_send<F: Future + Send>(future: F) -> F {
155            future
156        }
157
158        assert_eq!(assert_send(write_generic(Vec::new())).await?, vec![1, 2, 3]);
159        Ok(())
160    }
161
162    #[rstest]
163    #[case::single_write(vec![vec![1, 2, 3]], vec![1, 2, 3])]
164    #[case::two_writes(vec![vec![1, 2], vec![3, 4]], vec![1, 2, 3, 4])]
165    #[case::three_writes(vec![vec![1], vec![2], vec![3]], vec![1, 2, 3])]
166    #[case::with_empty(vec![vec![1, 2], vec![], vec![3, 4]], vec![1, 2, 3, 4])]
167    #[tokio::test]
168    async fn test_vec_multiple_writes(#[case] writes: Vec<Vec<u8>>, #[case] expected: Vec<u8>) {
169        let mut writer = Vec::new();
170
171        for data in writes {
172            VortexWrite::write_all(&mut writer, data).await.unwrap();
173        }
174
175        VortexWrite::flush(&mut writer).await.unwrap();
176        VortexWrite::shutdown(&mut writer).await.unwrap();
177        assert_eq!(writer, expected);
178    }
179
180    #[rstest]
181    #[case::single_write(vec![vec![5, 6, 7]], vec![5, 6, 7])]
182    #[case::two_writes(vec![vec![10], vec![20]], vec![10, 20])]
183    #[case::multiple_small(vec![vec![1], vec![2], vec![3], vec![4]], vec![1, 2, 3, 4])]
184    #[tokio::test]
185    async fn test_byte_buffer_mut_operations(
186        #[case] writes: Vec<Vec<u8>>,
187        #[case] expected: Vec<u8>,
188    ) {
189        let mut buffer = ByteBufferMut::with_capacity(0);
190
191        for data in writes {
192            VortexWrite::write_all(&mut buffer, data).await.unwrap();
193        }
194
195        VortexWrite::flush(&mut buffer).await.unwrap();
196        VortexWrite::shutdown(&mut buffer).await.unwrap();
197        assert_eq!(buffer.as_ref(), &expected[..]);
198    }
199
200    #[rstest]
201    #[case::empty(vec![], 0)]
202    #[case::single_byte(vec![42], 1)]
203    #[case::multiple_bytes(vec![1, 2, 3, 4, 5], 5)]
204    #[case::large(vec![0; 1024], 1024)]
205    #[tokio::test]
206    async fn test_various_write_sizes(#[case] data: Vec<u8>, #[case] expected_len: usize) {
207        let mut writer = Vec::new();
208        VortexWrite::write_all(&mut writer, data.clone())
209            .await
210            .unwrap();
211        assert_eq!(writer.len(), expected_len);
212        assert_eq!(writer, data);
213    }
214
215    #[tokio::test]
216    async fn test_cursor_operations() {
217        let mut data = [0u8; 20];
218        {
219            let mut cursor = Cursor::new(&mut data[..]);
220
221            // Write to cursor
222            VortexWrite::write_all(&mut cursor, vec![1, 2, 3, 4, 5])
223                .await
224                .unwrap();
225            assert_eq!(cursor.position(), 5);
226
227            // Write more data
228            VortexWrite::write_all(&mut cursor, vec![6, 7, 8, 9, 10])
229                .await
230                .unwrap();
231            assert_eq!(cursor.position(), 10);
232
233            // Test flush
234            VortexWrite::flush(&mut cursor).await.unwrap();
235        }
236
237        // Check data after cursor is dropped
238        assert_eq!(&data[..10], &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
239    }
240
241    #[tokio::test]
242    async fn test_futures_cursor() {
243        let mut vec = Vec::new();
244        {
245            let mut cursor = futures::io::Cursor::new(&mut vec);
246
247            // Test write operations
248            VortexWrite::write_all(&mut cursor, vec![10, 20, 30])
249                .await
250                .unwrap();
251            assert_eq!(cursor.position(), 3);
252
253            VortexWrite::write_all(&mut cursor, vec![40, 50])
254                .await
255                .unwrap();
256            assert_eq!(cursor.position(), 5);
257
258            // Test flush and shutdown
259            VortexWrite::flush(&mut cursor).await.unwrap();
260            VortexWrite::shutdown(&mut cursor).await.unwrap();
261        }
262
263        assert_eq!(vec, vec![10, 20, 30, 40, 50]);
264    }
265
266    #[tokio::test]
267    async fn test_mutable_reference() {
268        let mut vec = Vec::new();
269        {
270            let mut writer_ref = &mut vec;
271
272            // Test operations through mutable reference
273            VortexWrite::write_all(&mut writer_ref, vec![100, 101, 102])
274                .await
275                .unwrap();
276
277            VortexWrite::flush(&mut writer_ref).await.unwrap();
278            VortexWrite::shutdown(&mut writer_ref).await.unwrap();
279        }
280
281        assert_eq!(vec, vec![100, 101, 102]);
282    }
283
284    #[tokio::test]
285    async fn test_large_writes() {
286        let mut writer = Vec::new();
287        let large_data = vec![42u8; 100_000];
288
289        VortexWrite::write_all(&mut writer, large_data.clone())
290            .await
291            .unwrap();
292        assert_eq!(writer.len(), 100_000);
293        assert!(writer.iter().all(|&b| b == 42));
294    }
295
296    #[tokio::test]
297    async fn test_empty_writes() {
298        let mut writer = Vec::new();
299        let empty = vec![];
300
301        VortexWrite::write_all(&mut writer, empty.clone())
302            .await
303            .unwrap();
304        assert!(writer.is_empty());
305
306        VortexWrite::write_all(&mut writer, vec![1, 2, 3])
307            .await
308            .unwrap();
309        VortexWrite::write_all(&mut writer, empty).await.unwrap();
310        assert_eq!(writer, vec![1, 2, 3]);
311    }
312
313    #[tokio::test]
314    async fn test_sequential_accumulation() {
315        let mut buffer = ByteBufferMut::with_capacity(0);
316
317        for i in 0u8..5 {
318            VortexWrite::write_all(&mut buffer, vec![i]).await.unwrap();
319        }
320
321        assert_eq!(buffer.len(), 5);
322        assert_eq!(buffer.as_ref(), &[0, 1, 2, 3, 4]);
323    }
324
325    #[rstest]
326    #[case::vec_writer(0)]
327    #[case::byte_buffer(1)]
328    #[tokio::test]
329    async fn test_writer_types(#[case] writer_type: usize) {
330        let data = vec![10, 20, 30];
331
332        match writer_type {
333            0 => {
334                let mut writer = Vec::new();
335                VortexWrite::write_all(&mut writer, data.clone())
336                    .await
337                    .unwrap();
338                assert_eq!(writer, data);
339            }
340            1 => {
341                let mut writer = ByteBufferMut::with_capacity(0);
342                VortexWrite::write_all(&mut writer, data.clone())
343                    .await
344                    .unwrap();
345                assert_eq!(writer.as_ref(), &data[..]);
346            }
347            _ => unreachable!(),
348        }
349    }
350}
vortex_io/write.rs

vortex_io/
write.rs
