vortex_io/
write.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::future::{Future, ready};
5use std::io::{self, Cursor, Write};
6
7use vortex_buffer::ByteBufferMut;
8
9use crate::IoBuf;
10
11pub trait VortexWrite {
12    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>>;
13    fn flush(&mut self) -> impl Future<Output = io::Result<()>>;
14    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>>;
15}
16
17impl VortexWrite for Vec<u8> {
18    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> {
19        self.extend_from_slice(buffer.as_slice());
20        ready(Ok(buffer))
21    }
22
23    fn flush(&mut self) -> impl Future<Output = io::Result<()>> {
24        ready(Ok(()))
25    }
26
27    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> {
28        ready(Ok(()))
29    }
30}
31
32impl VortexWrite for ByteBufferMut {
33    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> {
34        self.extend_from_slice(buffer.as_slice());
35        ready(Ok(buffer))
36    }
37
38    fn flush(&mut self) -> impl Future<Output = io::Result<()>> {
39        ready(Ok(()))
40    }
41
42    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> {
43        ready(Ok(()))
44    }
45}
46
47impl<T> VortexWrite for Cursor<T>
48where
49    Cursor<T>: Write,
50{
51    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> {
52        ready(Write::write_all(self, buffer.as_slice()).map(|_| buffer))
53    }
54
55    fn flush(&mut self) -> impl Future<Output = io::Result<()>> {
56        ready(Write::flush(self))
57    }
58
59    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> {
60        ready(Ok(()))
61    }
62}
63
64impl<W: VortexWrite> VortexWrite for futures::io::Cursor<W> {
65    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> {
66        self.set_position(self.position() + buffer.as_slice().len() as u64);
67        VortexWrite::write_all(self.get_mut(), buffer)
68    }
69
70    fn flush(&mut self) -> impl Future<Output = io::Result<()>> {
71        VortexWrite::flush(self.get_mut())
72    }
73
74    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> {
75        VortexWrite::shutdown(self.get_mut())
76    }
77}
78
79impl<W: VortexWrite> VortexWrite for &mut W {
80    fn write_all<B: IoBuf>(&mut self, buffer: B) -> impl Future<Output = io::Result<B>> {
81        (*self).write_all(buffer)
82    }
83
84    fn flush(&mut self) -> impl Future<Output = io::Result<()>> {
85        (*self).flush()
86    }
87
88    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> {
89        (*self).shutdown()
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use rstest::rstest;
96    use vortex_buffer::ByteBufferMut;
97
98    use super::*;
99
100    #[rstest]
101    #[case::single_write(vec![vec![1, 2, 3]], vec![1, 2, 3])]
102    #[case::two_writes(vec![vec![1, 2], vec![3, 4]], vec![1, 2, 3, 4])]
103    #[case::three_writes(vec![vec![1], vec![2], vec![3]], vec![1, 2, 3])]
104    #[case::with_empty(vec![vec![1, 2], vec![], vec![3, 4]], vec![1, 2, 3, 4])]
105    #[tokio::test]
106    async fn test_vec_multiple_writes(#[case] writes: Vec<Vec<u8>>, #[case] expected: Vec<u8>) {
107        let mut writer = Vec::new();
108
109        for data in writes {
110            VortexWrite::write_all(&mut writer, data).await.unwrap();
111        }
112
113        VortexWrite::flush(&mut writer).await.unwrap();
114        VortexWrite::shutdown(&mut writer).await.unwrap();
115        assert_eq!(writer, expected);
116    }
117
118    #[rstest]
119    #[case::single_write(vec![vec![5, 6, 7]], vec![5, 6, 7])]
120    #[case::two_writes(vec![vec![10], vec![20]], vec![10, 20])]
121    #[case::multiple_small(vec![vec![1], vec![2], vec![3], vec![4]], vec![1, 2, 3, 4])]
122    #[tokio::test]
123    async fn test_byte_buffer_mut_operations(
124        #[case] writes: Vec<Vec<u8>>,
125        #[case] expected: Vec<u8>,
126    ) {
127        let mut buffer = ByteBufferMut::with_capacity(0);
128
129        for data in writes {
130            VortexWrite::write_all(&mut buffer, data).await.unwrap();
131        }
132
133        VortexWrite::flush(&mut buffer).await.unwrap();
134        VortexWrite::shutdown(&mut buffer).await.unwrap();
135        assert_eq!(buffer.as_ref(), &expected[..]);
136    }
137
138    #[rstest]
139    #[case::empty(vec![], 0)]
140    #[case::single_byte(vec![42], 1)]
141    #[case::multiple_bytes(vec![1, 2, 3, 4, 5], 5)]
142    #[case::large(vec![0; 1024], 1024)]
143    #[tokio::test]
144    async fn test_various_write_sizes(#[case] data: Vec<u8>, #[case] expected_len: usize) {
145        let mut writer = Vec::new();
146        VortexWrite::write_all(&mut writer, data.clone())
147            .await
148            .unwrap();
149        assert_eq!(writer.len(), expected_len);
150        assert_eq!(writer, data);
151    }
152
153    #[tokio::test]
154    async fn test_cursor_operations() {
155        let mut data = [0u8; 20];
156        {
157            let mut cursor = Cursor::new(&mut data[..]);
158
159            // Write to cursor
160            VortexWrite::write_all(&mut cursor, vec![1, 2, 3, 4, 5])
161                .await
162                .unwrap();
163            assert_eq!(cursor.position(), 5);
164
165            // Write more data
166            VortexWrite::write_all(&mut cursor, vec![6, 7, 8, 9, 10])
167                .await
168                .unwrap();
169            assert_eq!(cursor.position(), 10);
170
171            // Test flush
172            VortexWrite::flush(&mut cursor).await.unwrap();
173        }
174
175        // Check data after cursor is dropped
176        assert_eq!(&data[..10], &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
177    }
178
179    #[tokio::test]
180    async fn test_futures_cursor() {
181        let mut vec = Vec::new();
182        {
183            let mut cursor = futures::io::Cursor::new(&mut vec);
184
185            // Test write operations
186            VortexWrite::write_all(&mut cursor, vec![10, 20, 30])
187                .await
188                .unwrap();
189            assert_eq!(cursor.position(), 3);
190
191            VortexWrite::write_all(&mut cursor, vec![40, 50])
192                .await
193                .unwrap();
194            assert_eq!(cursor.position(), 5);
195
196            // Test flush and shutdown
197            VortexWrite::flush(&mut cursor).await.unwrap();
198            VortexWrite::shutdown(&mut cursor).await.unwrap();
199        }
200
201        assert_eq!(vec, vec![10, 20, 30, 40, 50]);
202    }
203
204    #[tokio::test]
205    async fn test_mutable_reference() {
206        let mut vec = Vec::new();
207        {
208            let mut writer_ref = &mut vec;
209
210            // Test operations through mutable reference
211            VortexWrite::write_all(&mut writer_ref, vec![100, 101, 102])
212                .await
213                .unwrap();
214
215            VortexWrite::flush(&mut writer_ref).await.unwrap();
216            VortexWrite::shutdown(&mut writer_ref).await.unwrap();
217        }
218
219        assert_eq!(vec, vec![100, 101, 102]);
220    }
221
222    #[tokio::test]
223    async fn test_large_writes() {
224        let mut writer = Vec::new();
225        let large_data = vec![42u8; 100_000];
226
227        VortexWrite::write_all(&mut writer, large_data.clone())
228            .await
229            .unwrap();
230        assert_eq!(writer.len(), 100_000);
231        assert!(writer.iter().all(|&b| b == 42));
232    }
233
234    #[tokio::test]
235    async fn test_empty_writes() {
236        let mut writer = Vec::new();
237        let empty = vec![];
238
239        VortexWrite::write_all(&mut writer, empty.clone())
240            .await
241            .unwrap();
242        assert!(writer.is_empty());
243
244        VortexWrite::write_all(&mut writer, vec![1, 2, 3])
245            .await
246            .unwrap();
247        VortexWrite::write_all(&mut writer, empty).await.unwrap();
248        assert_eq!(writer, vec![1, 2, 3]);
249    }
250
251    #[tokio::test]
252    async fn test_sequential_accumulation() {
253        let mut buffer = ByteBufferMut::with_capacity(0);
254
255        for i in 0u8..5 {
256            VortexWrite::write_all(&mut buffer, vec![i]).await.unwrap();
257        }
258
259        assert_eq!(buffer.len(), 5);
260        assert_eq!(buffer.as_ref(), &[0, 1, 2, 3, 4]);
261    }
262
263    #[rstest]
264    #[case::vec_writer(0)]
265    #[case::byte_buffer(1)]
266    #[tokio::test]
267    async fn test_writer_types(#[case] writer_type: usize) {
268        let data = vec![10, 20, 30];
269
270        match writer_type {
271            0 => {
272                let mut writer = Vec::new();
273                VortexWrite::write_all(&mut writer, data.clone())
274                    .await
275                    .unwrap();
276                assert_eq!(writer, data);
277            }
278            1 => {
279                let mut writer = ByteBufferMut::with_capacity(0);
280                VortexWrite::write_all(&mut writer, data.clone())
281                    .await
282                    .unwrap();
283                assert_eq!(writer.as_ref(), &data[..]);
284            }
285            _ => unreachable!(),
286        }
287    }
288}
vortex_io/write.rs

vortex_io/
write.rs
