A better Rust ATProto crate
1use crate::io::{self, ErrorKind, Initializer, Read, Write};
2use core::{fmt, mem::MaybeUninit};
3/// Copies the entire contents of a reader into a writer.
4///
5/// This function will continuously read data from `reader` and then
6/// write it into `writer` in a streaming fashion until `reader`
7/// returns EOF.
8///
9/// On success, the total number of bytes that were copied from
10/// `reader` to `writer` is returned.
11///
12/// # Errors
13///
14/// This function will return an error immediately if any call to `read` or
15/// `write` returns an error. All instances of `ErrorKind::Interrupted` are
16/// handled by this function and the underlying operation is retried.
17pub fn copy<R, W>(reader: &mut R, writer: &mut W) -> io::Result<u64>
18where
19 R: ?Sized + Read,
20 W: ?Sized + Write,
21{
22 let mut buf = MaybeUninit::<[u8; super::DEFAULT_BUF_SIZE]>::uninit();
23 // FIXME(eliza): the stdlib has a scary comment that says we haven't decided
24 // if this is okay or not...
25 unsafe {
26 reader.initializer().initialize(&mut *buf.as_mut_ptr());
27 }
28
29 let mut written = 0;
30 loop {
31 let len = match reader.read(unsafe { &mut *buf.as_mut_ptr() }) {
32 Ok(0) => return Ok(written),
33 Ok(len) => len,
34 Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
35 Err(e) => return Err(e),
36 };
37 writer.write_all(unsafe { &(&*buf.as_ptr())[..len] })?;
38 written += len as u64;
39 }
40}
41
42/// A reader which is always at EOF.
43///
44/// This struct is generally created by calling [`empty`]. Please see
45/// the documentation of [`empty()`][`empty`] for more details.
46///
47/// [`empty`]: fn.empty.html
48pub struct Empty {
49 _priv: (),
50}
51
52/// Constructs a new handle to an empty reader.
53///
54/// All reads from the returned reader will return `Ok(0)`.
55///
56pub fn empty() -> Empty {
57 Empty { _priv: () }
58}
59
60impl Read for Empty {
61 #[inline]
62 fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> {
63 Ok(0)
64 }
65
66 #[inline]
67 unsafe fn initializer(&self) -> Initializer {
68 unsafe { Initializer::nop() }
69 }
70}
71
72// impl BufRead for Empty {
73// #[inline]
74// fn fill_buf(&mut self) -> io::Result<&[u8]> {
75// Ok(&[])
76// }
77// #[inline]
78// fn consume(&mut self, _n: usize) {}
79// }
80
81impl fmt::Debug for Empty {
82 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
83 f.pad("Empty { .. }")
84 }
85}
86
87/// A reader which yields one byte over and over and over and over and over and...
88///
89/// This struct is generally created by calling [`repeat`][repeat]. Please
90/// see the documentation of `repeat()` for more details.
91///
92/// [repeat]: fn.repeat.html
93pub struct Repeat {
94 byte: u8,
95}
96
97/// Creates an instance of a reader that infinitely repeats one byte.
98///
99/// All reads from this reader will succeed by filling the specified buffer with
100/// the given byte.
101pub fn repeat(byte: u8) -> Repeat {
102 Repeat { byte }
103}
104
105impl Read for Repeat {
106 #[inline]
107 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
108 for slot in &mut *buf {
109 *slot = self.byte;
110 }
111 Ok(buf.len())
112 }
113
114 #[inline]
115 unsafe fn initializer(&self) -> Initializer {
116 unsafe { Initializer::nop() }
117 }
118}
119
120impl fmt::Debug for Repeat {
121 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
122 f.pad("Repeat { .. }")
123 }
124}
125
126/// A writer which will move data into the void.
127///
128/// This struct is generally created by calling [`sink`][sink]. Please
129/// see the documentation of `sink()` for more details.
130///
131/// [sink]: fn.sink.html
132pub struct Sink {
133 _priv: (),
134}
135
136/// Creates an instance of a writer which will successfully consume all data.
137///
138/// All calls to `write` on the returned instance will return `Ok(buf.len())`
139/// and the contents of the buffer will not be inspected.
140pub fn sink() -> Sink {
141 Sink { _priv: () }
142}
143
144impl Write for Sink {
145 #[inline]
146 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
147 Ok(buf.len())
148 }
149
150 #[inline]
151 fn flush(&mut self) -> io::Result<()> {
152 Ok(())
153 }
154}
155
156impl fmt::Debug for Sink {
157 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158 f.pad("Sink { .. }")
159 }
160}
161
162#[cfg(all(test, not(loom)))]
163mod tests {
164 use crate::io::prelude::*;
165 use crate::io::{copy, empty, repeat, sink};
166
167 #[test]
168 fn copy_copies() {
169 let mut r = repeat(0).take(4);
170 let mut w = sink();
171 assert_eq!(copy(&mut r, &mut w).unwrap(), 4);
172
173 let mut r = repeat(0).take(1 << 17);
174 assert_eq!(
175 copy(&mut r as &mut dyn Read, &mut w as &mut dyn Write).unwrap(),
176 1 << 17
177 );
178 }
179
180 #[test]
181 fn sink_sinks() {
182 let mut s = sink();
183 assert_eq!(s.write(&[]).unwrap(), 0);
184 assert_eq!(s.write(&[0]).unwrap(), 1);
185 assert_eq!(s.write(&[0; 1024]).unwrap(), 1024);
186 assert_eq!(s.by_ref().write(&[0; 1024]).unwrap(), 1024);
187 }
188
189 #[test]
190 fn empty_reads() {
191 let mut e = empty();
192 assert_eq!(e.read(&mut []).unwrap(), 0);
193 assert_eq!(e.read(&mut [0]).unwrap(), 0);
194 assert_eq!(e.read(&mut [0; 1024]).unwrap(), 0);
195 assert_eq!(e.by_ref().read(&mut [0; 1024]).unwrap(), 0);
196 }
197
198 #[test]
199 fn repeat_repeats() {
200 let mut r = repeat(4);
201 let mut b = [0; 1024];
202 assert_eq!(r.read(&mut b).unwrap(), 1024);
203 assert!(b.iter().all(|b| *b == 4));
204 }
205
206 #[test]
207 fn take_some_bytes() {
208 assert_eq!(repeat(4).take(100).bytes().count(), 100);
209 assert_eq!(repeat(4).take(100).bytes().next().unwrap().unwrap(), 4);
210 assert_eq!(
211 repeat(1).take(10).chain(repeat(2).take(10)).bytes().count(),
212 20
213 );
214 }
215}