1 use crate::io::util::read_until::read_until_internal;
2 use crate::io::AsyncBufRead;
3
4 use std::future::Future;
5 use std::io;
6 use std::mem;
7 use std::pin::Pin;
8 use std::string::FromUtf8Error;
9 use std::task::{Context, Poll};
10
11 cfg_io_util! {
12 /// Future for the [`read_line`](crate::io::AsyncBufReadExt::read_line) method.
13 #[derive(Debug)]
14 #[must_use = "futures do nothing unless you `.await` or poll them"]
15 pub struct ReadLine<'a, R: ?Sized> {
16 reader: &'a mut R,
17 /// This is the buffer we were provided. It will be replaced with an empty string
18 /// while reading to postpone utf-8 handling until after reading.
19 output: &'a mut String,
20 /// The actual allocation of the string is moved into this vector instead.
21 buf: Vec<u8>,
22 /// The number of bytes appended to buf. This can be less than buf.len() if
23 /// the buffer was not empty when the operation was started.
24 read: usize,
25 }
26 }
27
read_line<'a, R>(reader: &'a mut R, string: &'a mut String) -> ReadLine<'a, R> where R: AsyncBufRead + ?Sized + Unpin,28 pub(crate) fn read_line<'a, R>(reader: &'a mut R, string: &'a mut String) -> ReadLine<'a, R>
29 where
30 R: AsyncBufRead + ?Sized + Unpin,
31 {
32 ReadLine {
33 reader,
34 buf: mem::replace(string, String::new()).into_bytes(),
35 output: string,
36 read: 0,
37 }
38 }
39
put_back_original_data(output: &mut String, mut vector: Vec<u8>, num_bytes_read: usize)40 fn put_back_original_data(output: &mut String, mut vector: Vec<u8>, num_bytes_read: usize) {
41 let original_len = vector.len() - num_bytes_read;
42 vector.truncate(original_len);
43 *output = String::from_utf8(vector).expect("The original data must be valid utf-8.");
44 }
45
46 /// This handles the various failure cases and puts the string back into `output`.
47 ///
48 /// The `truncate_on_io_error` bool is necessary because `read_to_string` and `read_line`
49 /// disagree on what should happen when an IO error occurs.
finish_string_read( io_res: io::Result<usize>, utf8_res: Result<String, FromUtf8Error>, read: usize, output: &mut String, truncate_on_io_error: bool, ) -> Poll<io::Result<usize>>50 pub(super) fn finish_string_read(
51 io_res: io::Result<usize>,
52 utf8_res: Result<String, FromUtf8Error>,
53 read: usize,
54 output: &mut String,
55 truncate_on_io_error: bool,
56 ) -> Poll<io::Result<usize>> {
57 match (io_res, utf8_res) {
58 (Ok(num_bytes), Ok(string)) => {
59 debug_assert_eq!(read, 0);
60 *output = string;
61 Poll::Ready(Ok(num_bytes))
62 }
63 (Err(io_err), Ok(string)) => {
64 *output = string;
65 if truncate_on_io_error {
66 let original_len = output.len() - read;
67 output.truncate(original_len);
68 }
69 Poll::Ready(Err(io_err))
70 }
71 (Ok(num_bytes), Err(utf8_err)) => {
72 debug_assert_eq!(read, 0);
73 put_back_original_data(output, utf8_err.into_bytes(), num_bytes);
74
75 Poll::Ready(Err(io::Error::new(
76 io::ErrorKind::InvalidData,
77 "stream did not contain valid UTF-8",
78 )))
79 }
80 (Err(io_err), Err(utf8_err)) => {
81 put_back_original_data(output, utf8_err.into_bytes(), read);
82
83 Poll::Ready(Err(io_err))
84 }
85 }
86 }
87
read_line_internal<R: AsyncBufRead + ?Sized>( reader: Pin<&mut R>, cx: &mut Context<'_>, output: &mut String, buf: &mut Vec<u8>, read: &mut usize, ) -> Poll<io::Result<usize>>88 pub(super) fn read_line_internal<R: AsyncBufRead + ?Sized>(
89 reader: Pin<&mut R>,
90 cx: &mut Context<'_>,
91 output: &mut String,
92 buf: &mut Vec<u8>,
93 read: &mut usize,
94 ) -> Poll<io::Result<usize>> {
95 let io_res = ready!(read_until_internal(reader, cx, b'\n', buf, read));
96 let utf8_res = String::from_utf8(mem::replace(buf, Vec::new()));
97
98 // At this point both buf and output are empty. The allocation is in utf8_res.
99
100 debug_assert!(buf.is_empty());
101 debug_assert!(output.is_empty());
102 finish_string_read(io_res, utf8_res, *read, output, false)
103 }
104
105 impl<R: AsyncBufRead + ?Sized + Unpin> Future for ReadLine<'_, R> {
106 type Output = io::Result<usize>;
107
poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>108 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
109 let Self {
110 reader,
111 output,
112 buf,
113 read,
114 } = &mut *self;
115
116 read_line_internal(Pin::new(reader), cx, output, buf, read)
117 }
118 }
119
120 #[cfg(test)]
121 mod tests {
122 use super::*;
123
124 #[test]
assert_unpin()125 fn assert_unpin() {
126 use std::marker::PhantomPinned;
127 crate::is_unpin::<ReadLine<'_, PhantomPinned>>();
128 }
129 }
130