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