1use std::{
2 io::{self, IsTerminal, Read, Write},
3 os::windows::io::{AsRawHandle, BorrowedHandle, RawHandle},
4 sync::OnceLock,
5 task::Poll,
6};
7
8use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut};
9use compio_driver::{
10 AsFd, AsRawFd, BorrowedFd, BufferRef, OpCode, OpType, RawFd, ResultTakeBuffer, SharedFd,
11 op::{BufResultExt, Read as OpRead, ReadManaged, Write as OpWrite},
12};
13use compio_io::{AsyncRead, AsyncReadManaged, AsyncReadMulti, AsyncWrite};
14use compio_runtime::Runtime;
15use futures_util::{Stream, StreamExt};
16use windows_sys::Win32::System::IO::OVERLAPPED;
17
18#[cfg(doc)]
19use super::{stderr, stdin, stdout};
20
21struct StdRead<R: Read, B: IoBufMut> {
22 reader: R,
23 buffer: B,
24}
25
26impl<R: Read, B: IoBufMut> StdRead<R, B> {
27 pub fn new(reader: R, buffer: B) -> Self {
28 Self { reader, buffer }
29 }
30}
31
32unsafe impl<R: Read, B: IoBufMut> OpCode for StdRead<R, B> {
33 type Control = ();
34
35 unsafe fn init(&mut self, _: &mut Self::Control) {}
36
37 fn op_type(&self, _: &Self::Control) -> OpType {
38 OpType::Blocking
39 }
40
41 unsafe fn operate(
42 &mut self,
43 _: &mut Self::Control,
44 _optr: *mut OVERLAPPED,
45 ) -> Poll<io::Result<usize>> {
46 #[cfg(feature = "read_buf")]
47 {
48 let slice = self.buffer.as_uninit();
49 let mut buf = io::BorrowedBuf::from(slice);
50 let mut cursor = buf.unfilled();
51 self.reader.read_buf(cursor.reborrow())?;
52 Poll::Ready(Ok(cursor.written()))
53 }
54 #[cfg(not(feature = "read_buf"))]
55 {
56 use compio_buf::IoBufMutExt;
57
58 let slice = self.buffer.ensure_init();
59 self.reader.read(slice).into()
60 }
61 }
62}
63
64impl<R: Read, B: IoBufMut> IntoInner for StdRead<R, B> {
65 type Inner = B;
66
67 fn into_inner(self) -> Self::Inner {
68 self.buffer
69 }
70}
71
72struct StdWrite<W: Write, B: IoBuf> {
73 writer: W,
74 buffer: B,
75}
76
77impl<W: Write, B: IoBuf> StdWrite<W, B> {
78 pub fn new(writer: W, buffer: B) -> Self {
79 Self { writer, buffer }
80 }
81}
82
83unsafe impl<W: Write, B: IoBuf> OpCode for StdWrite<W, B> {
84 type Control = ();
85
86 unsafe fn init(&mut self, _: &mut Self::Control) {}
87
88 fn op_type(&self, _: &Self::Control) -> OpType {
89 OpType::Blocking
90 }
91
92 unsafe fn operate(
93 &mut self,
94 _: &mut Self::Control,
95 _optr: *mut OVERLAPPED,
96 ) -> Poll<io::Result<usize>> {
97 let slice = self.buffer.as_init();
98 self.writer.write(slice).into()
99 }
100}
101
102impl<W: Write, B: IoBuf> IntoInner for StdWrite<W, B> {
103 type Inner = B;
104
105 fn into_inner(self) -> Self::Inner {
106 self.buffer
107 }
108}
109
110#[derive(Debug)]
111struct StaticFd(RawHandle);
112
113impl AsFd for StaticFd {
114 fn as_fd(&self) -> BorrowedFd<'_> {
115 BorrowedFd::File(unsafe { BorrowedHandle::borrow_raw(self.0) })
117 }
118}
119
120impl AsRawFd for StaticFd {
121 fn as_raw_fd(&self) -> RawFd {
122 self.0 as _
123 }
124}
125
126static STDIN_ISATTY: OnceLock<bool> = OnceLock::new();
127
128#[derive(Debug, Clone)]
132pub struct Stdin {
133 fd: SharedFd<StaticFd>,
134 isatty: bool,
135}
136
137impl Stdin {
138 pub(crate) fn new() -> Self {
139 let stdin = io::stdin();
140 let isatty = *STDIN_ISATTY.get_or_init(|| {
141 stdin.is_terminal()
142 || Runtime::with_current(|r| r.attach(stdin.as_raw_handle() as _)).is_err()
143 });
144 Self {
145 fd: SharedFd::new(StaticFd(stdin.as_raw_handle())),
146 isatty,
147 }
148 }
149}
150
151impl AsyncRead for Stdin {
152 async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
153 let res = if self.isatty {
154 let op = StdRead::new(io::stdin(), buf);
155 compio_runtime::submit(op).await.into_inner()
156 } else {
157 let op = OpRead::new(self.fd.clone(), buf);
158 compio_runtime::submit(op).await.into_inner()
159 };
160 unsafe { res.map_advanced() }
161 }
162}
163
164impl AsyncReadManaged for Stdin {
165 type Buffer = BufferRef;
166
167 async fn read_managed(&mut self, len: usize) -> io::Result<Option<Self::Buffer>> {
168 (&*self).read_managed(len).await
169 }
170}
171
172impl AsyncReadManaged for &Stdin {
173 type Buffer = BufferRef;
174
175 async fn read_managed(&mut self, len: usize) -> io::Result<Option<Self::Buffer>> {
176 let runtime = Runtime::current();
177 let buffer_pool = runtime.buffer_pool()?;
178 if self.isatty {
179 let buf = buffer_pool.pop()?;
180 let op = StdRead::new(io::stdin(), buf);
181 unsafe { compio_runtime::submit(op).await.take_buffer() }
182 } else {
183 let op = ReadManaged::new(self.fd.clone(), &buffer_pool, len)?;
184 unsafe { compio_runtime::submit(op).await.take_buffer() }
185 }
186 }
187}
188
189impl AsyncReadMulti for Stdin {
190 fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
191 futures_util::stream::once(self.read_managed(len))
192 .filter_map(|res| std::future::ready(res.transpose()))
193 }
194}
195
196impl AsyncReadMulti for &Stdin {
197 fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
198 futures_util::stream::once(self.read_managed(len))
199 .filter_map(|res| std::future::ready(res.transpose()))
200 }
201}
202
203impl AsRawFd for Stdin {
204 fn as_raw_fd(&self) -> RawFd {
205 self.fd.as_raw_fd()
206 }
207}
208
209static STDOUT_ISATTY: OnceLock<bool> = OnceLock::new();
210
211#[derive(Debug, Clone)]
215pub struct Stdout {
216 fd: SharedFd<StaticFd>,
217 isatty: bool,
218}
219
220impl Stdout {
221 pub(crate) fn new() -> Self {
222 let stdout = io::stdout();
223 let isatty = *STDOUT_ISATTY.get_or_init(|| {
224 stdout.is_terminal()
225 || Runtime::with_current(|r| r.attach(stdout.as_raw_handle() as _)).is_err()
226 });
227 Self {
228 fd: SharedFd::new(StaticFd(stdout.as_raw_handle())),
229 isatty,
230 }
231 }
232}
233
234impl AsyncWrite for Stdout {
235 async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
236 if self.isatty {
237 let op = StdWrite::new(io::stdout(), buf);
238 compio_runtime::submit(op).await.into_inner()
239 } else {
240 let op = OpWrite::new(self.fd.clone(), buf);
241 compio_runtime::submit(op).await.into_inner()
242 }
243 }
244
245 async fn flush(&mut self) -> io::Result<()> {
246 Ok(())
247 }
248
249 async fn shutdown(&mut self) -> io::Result<()> {
250 self.flush().await
251 }
252}
253
254impl AsRawFd for Stdout {
255 fn as_raw_fd(&self) -> RawFd {
256 self.fd.as_raw_fd()
257 }
258}
259
260static STDERR_ISATTY: OnceLock<bool> = OnceLock::new();
261
262#[derive(Debug, Clone)]
266pub struct Stderr {
267 fd: SharedFd<StaticFd>,
268 isatty: bool,
269}
270
271impl Stderr {
272 pub(crate) fn new() -> Self {
273 let stderr = io::stderr();
274 let isatty = *STDERR_ISATTY.get_or_init(|| {
275 stderr.is_terminal()
276 || Runtime::with_current(|r| r.attach(stderr.as_raw_handle() as _)).is_err()
277 });
278 Self {
279 fd: SharedFd::new(StaticFd(stderr.as_raw_handle())),
280 isatty,
281 }
282 }
283}
284
285impl AsyncWrite for Stderr {
286 async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
287 if self.isatty {
288 let op = StdWrite::new(io::stderr(), buf);
289 compio_runtime::submit(op).await.into_inner()
290 } else {
291 let op = OpWrite::new(self.fd.clone(), buf);
292 compio_runtime::submit(op).await.into_inner()
293 }
294 }
295
296 async fn flush(&mut self) -> io::Result<()> {
297 Ok(())
298 }
299
300 async fn shutdown(&mut self) -> io::Result<()> {
301 self.flush().await
302 }
303}
304
305impl AsRawFd for Stderr {
306 fn as_raw_fd(&self) -> RawFd {
307 self.fd.as_raw_fd()
308 }
309}