|
261 | 261 |
|
262 | 262 | use crate::cmp;
|
263 | 263 | use crate::fmt;
|
| 264 | +use crate::mem; |
264 | 265 | use crate::memchr;
|
265 | 266 | use crate::ops::{Deref, DerefMut};
|
266 | 267 | use crate::ptr;
|
@@ -1376,6 +1377,70 @@ pub trait Write {
|
1376 | 1377 | Ok(())
|
1377 | 1378 | }
|
1378 | 1379 |
|
| 1380 | + /// Attempts to write multiple buffers into this writer. |
| 1381 | + /// |
| 1382 | + /// This method will continuously call [`write_vectored`] until there is no |
| 1383 | + /// more data to be written or an error of non-[`ErrorKind::Interrupted`] |
| 1384 | + /// kind is returned. This method will not return until all buffers have |
| 1385 | + /// been successfully written or such an error occurs. The first error that |
| 1386 | + /// is not of [`ErrorKind::Interrupted`] kind generated from this method |
| 1387 | + /// will be returned. |
| 1388 | + /// |
| 1389 | + /// If the buffer contains no data, this will never call [`write_vectored`]. |
| 1390 | + /// |
| 1391 | + /// [`write_vectored`]: #method.write_vectored |
| 1392 | + /// [`ErrorKind::Interrupted`]: ../../std/io/enum.ErrorKind.html#variant.Interrupted |
| 1393 | + /// |
| 1394 | + /// # Notes |
| 1395 | + /// |
| 1396 | + /// |
| 1397 | + /// Unlike `io::Write::write_vectored`, this takes a *mutable* reference to |
| 1398 | + /// a slice of `IoSlice`s, not an immutable one. That's because we need to |
| 1399 | + /// modify the slice to keep track of the bytes already written. |
| 1400 | + /// |
| 1401 | + /// Once this function returns, the contents of `bufs` are unspecified, as |
| 1402 | + /// this depends on how many calls to `write_vectored` were necessary. It is |
| 1403 | + /// best to understand this function as taking ownership of `bufs` and to |
| 1404 | + /// not use `bufs` afterwards. The underlying buffers, to which the |
| 1405 | + /// `IoSlice`s point (but not the `IoSlice`s themselves), are unchanged and |
| 1406 | + /// can be reused. |
| 1407 | + /// |
| 1408 | + /// # Examples |
| 1409 | + /// |
| 1410 | + /// ``` |
| 1411 | + /// #![feature(write_all_vectored)] |
| 1412 | + /// # fn main() -> std::io::Result<()> { |
| 1413 | + /// |
| 1414 | + /// use std::io::{Write, IoSlice}; |
| 1415 | + /// |
| 1416 | + /// let mut writer = Vec::new(); |
| 1417 | + /// let bufs = &mut [ |
| 1418 | + /// IoSlice::new(&[1]), |
| 1419 | + /// IoSlice::new(&[2, 3]), |
| 1420 | + /// IoSlice::new(&[4, 5, 6]), |
| 1421 | + /// ]; |
| 1422 | + /// |
| 1423 | + /// writer.write_all_vectored(bufs)?; |
| 1424 | + /// // Note: the contents of `bufs` is now undefined, see the Notes section. |
| 1425 | + /// |
| 1426 | + /// assert_eq!(writer, &[1, 2, 3, 4, 5, 6]); |
| 1427 | + /// # Ok(()) } |
| 1428 | + /// ``` |
| 1429 | + #[unstable(feature = "write_all_vectored", issue = "70436")] |
| 1430 | + fn write_all_vectored(&mut self, mut bufs: &mut [IoSlice<'_>]) -> Result<()> { |
| 1431 | + while !bufs.is_empty() { |
| 1432 | + match self.write_vectored(bufs) { |
| 1433 | + Ok(0) => { |
| 1434 | + return Err(Error::new(ErrorKind::WriteZero, "failed to write whole buffer")); |
| 1435 | + } |
| 1436 | + Ok(n) => bufs = IoSlice::advance(mem::take(&mut bufs), n), |
| 1437 | + Err(ref e) if e.kind() == ErrorKind::Interrupted => {} |
| 1438 | + Err(e) => return Err(e), |
| 1439 | + } |
| 1440 | + } |
| 1441 | + Ok(()) |
| 1442 | + } |
| 1443 | + |
1379 | 1444 | /// Writes a formatted string into this writer, returning any error
|
1380 | 1445 | /// encountered.
|
1381 | 1446 | ///
|
@@ -2423,7 +2488,7 @@ impl<B: BufRead> Iterator for Lines<B> {
|
2423 | 2488 | #[cfg(test)]
|
2424 | 2489 | mod tests {
|
2425 | 2490 | use super::{repeat, Cursor, SeekFrom};
|
2426 |
| - use crate::cmp; |
| 2491 | + use crate::cmp::{self, min}; |
2427 | 2492 | use crate::io::prelude::*;
|
2428 | 2493 | use crate::io::{self, IoSlice, IoSliceMut};
|
2429 | 2494 | use crate::ops::Deref;
|
@@ -2812,4 +2877,107 @@ mod tests {
|
2812 | 2877 | bufs = IoSlice::advance(bufs, 9);
|
2813 | 2878 | assert!(bufs.is_empty());
|
2814 | 2879 | }
|
| 2880 | + |
| 2881 | + /// Create a new writer that reads from at most `n_bufs` and reads |
| 2882 | + /// `per_call` bytes (in total) per call to write. |
| 2883 | + fn test_writer(n_bufs: usize, per_call: usize) -> TestWriter { |
| 2884 | + TestWriter { n_bufs, per_call, written: Vec::new() } |
| 2885 | + } |
| 2886 | + |
| 2887 | + struct TestWriter { |
| 2888 | + n_bufs: usize, |
| 2889 | + per_call: usize, |
| 2890 | + written: Vec<u8>, |
| 2891 | + } |
| 2892 | + |
| 2893 | + impl Write for TestWriter { |
| 2894 | + fn write(&mut self, buf: &[u8]) -> io::Result<usize> { |
| 2895 | + self.write_vectored(&[IoSlice::new(buf)]) |
| 2896 | + } |
| 2897 | + |
| 2898 | + fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> { |
| 2899 | + let mut left = self.per_call; |
| 2900 | + let mut written = 0; |
| 2901 | + for buf in bufs.iter().take(self.n_bufs) { |
| 2902 | + let n = min(left, buf.len()); |
| 2903 | + self.written.extend_from_slice(&buf[0..n]); |
| 2904 | + left -= n; |
| 2905 | + written += n; |
| 2906 | + } |
| 2907 | + Ok(written) |
| 2908 | + } |
| 2909 | + |
| 2910 | + fn flush(&mut self) -> io::Result<()> { |
| 2911 | + Ok(()) |
| 2912 | + } |
| 2913 | + } |
| 2914 | + |
| 2915 | + #[test] |
| 2916 | + fn test_writer_read_from_one_buf() { |
| 2917 | + let mut writer = test_writer(1, 2); |
| 2918 | + |
| 2919 | + assert_eq!(writer.write(&[]).unwrap(), 0); |
| 2920 | + assert_eq!(writer.write_vectored(&[]).unwrap(), 0); |
| 2921 | + |
| 2922 | + // Read at most 2 bytes. |
| 2923 | + assert_eq!(writer.write(&[1, 1, 1]).unwrap(), 2); |
| 2924 | + let bufs = &[IoSlice::new(&[2, 2, 2])]; |
| 2925 | + assert_eq!(writer.write_vectored(bufs).unwrap(), 2); |
| 2926 | + |
| 2927 | + // Only read from first buf. |
| 2928 | + let bufs = &[IoSlice::new(&[3]), IoSlice::new(&[4, 4])]; |
| 2929 | + assert_eq!(writer.write_vectored(bufs).unwrap(), 1); |
| 2930 | + |
| 2931 | + assert_eq!(writer.written, &[1, 1, 2, 2, 3]); |
| 2932 | + } |
| 2933 | + |
| 2934 | + #[test] |
| 2935 | + fn test_writer_read_from_multiple_bufs() { |
| 2936 | + let mut writer = test_writer(3, 3); |
| 2937 | + |
| 2938 | + // Read at most 3 bytes from two buffers. |
| 2939 | + let bufs = &[IoSlice::new(&[1]), IoSlice::new(&[2, 2, 2])]; |
| 2940 | + assert_eq!(writer.write_vectored(bufs).unwrap(), 3); |
| 2941 | + |
| 2942 | + // Read at most 3 bytes from three buffers. |
| 2943 | + let bufs = &[IoSlice::new(&[3]), IoSlice::new(&[4]), IoSlice::new(&[5, 5])]; |
| 2944 | + assert_eq!(writer.write_vectored(bufs).unwrap(), 3); |
| 2945 | + |
| 2946 | + assert_eq!(writer.written, &[1, 2, 2, 3, 4, 5]); |
| 2947 | + } |
| 2948 | + |
| 2949 | + #[test] |
| 2950 | + fn test_write_all_vectored() { |
| 2951 | + #[rustfmt::skip] // Becomes unreadable otherwise. |
| 2952 | + let tests: Vec<(_, &'static [u8])> = vec![ |
| 2953 | + (vec![], &[]), |
| 2954 | + (vec![IoSlice::new(&[1])], &[1]), |
| 2955 | + (vec![IoSlice::new(&[1, 2])], &[1, 2]), |
| 2956 | + (vec![IoSlice::new(&[1, 2, 3])], &[1, 2, 3]), |
| 2957 | + (vec![IoSlice::new(&[1, 2, 3, 4])], &[1, 2, 3, 4]), |
| 2958 | + (vec![IoSlice::new(&[1, 2, 3, 4, 5])], &[1, 2, 3, 4, 5]), |
| 2959 | + (vec![IoSlice::new(&[1]), IoSlice::new(&[2])], &[1, 2]), |
| 2960 | + (vec![IoSlice::new(&[1]), IoSlice::new(&[2, 2])], &[1, 2, 2]), |
| 2961 | + (vec![IoSlice::new(&[1, 1]), IoSlice::new(&[2, 2])], &[1, 1, 2, 2]), |
| 2962 | + (vec![IoSlice::new(&[1, 1]), IoSlice::new(&[2, 2, 2])], &[1, 1, 2, 2, 2]), |
| 2963 | + (vec![IoSlice::new(&[1, 1]), IoSlice::new(&[2, 2, 2])], &[1, 1, 2, 2, 2]), |
| 2964 | + (vec![IoSlice::new(&[1, 1, 1]), IoSlice::new(&[2, 2, 2])], &[1, 1, 1, 2, 2, 2]), |
| 2965 | + (vec![IoSlice::new(&[1, 1, 1]), IoSlice::new(&[2, 2, 2, 2])], &[1, 1, 1, 2, 2, 2, 2]), |
| 2966 | + (vec![IoSlice::new(&[1, 1, 1, 1]), IoSlice::new(&[2, 2, 2, 2])], &[1, 1, 1, 1, 2, 2, 2, 2]), |
| 2967 | + (vec![IoSlice::new(&[1]), IoSlice::new(&[2]), IoSlice::new(&[3])], &[1, 2, 3]), |
| 2968 | + (vec![IoSlice::new(&[1, 1]), IoSlice::new(&[2, 2]), IoSlice::new(&[3, 3])], &[1, 1, 2, 2, 3, 3]), |
| 2969 | + (vec![IoSlice::new(&[1]), IoSlice::new(&[2, 2]), IoSlice::new(&[3, 3, 3])], &[1, 2, 2, 3, 3, 3]), |
| 2970 | + (vec![IoSlice::new(&[1, 1, 1]), IoSlice::new(&[2, 2, 2]), IoSlice::new(&[3, 3, 3])], &[1, 1, 1, 2, 2, 2, 3, 3, 3]), |
| 2971 | + ]; |
| 2972 | + |
| 2973 | + let writer_configs = &[(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]; |
| 2974 | + |
| 2975 | + for (n_bufs, per_call) in writer_configs.iter().copied() { |
| 2976 | + for (mut input, wanted) in tests.clone().into_iter() { |
| 2977 | + let mut writer = test_writer(n_bufs, per_call); |
| 2978 | + assert!(writer.write_all_vectored(&mut *input).is_ok()); |
| 2979 | + assert_eq!(&*writer.written, &*wanted); |
| 2980 | + } |
| 2981 | + } |
| 2982 | + } |
2815 | 2983 | }
|
0 commit comments