use std::future::Future;
use std::pin::Pin;
use std::sync::Mutex;
use crate::future;
use crate::io::{self, Read};
use crate::task::{spawn_blocking, Context, JoinHandle, Poll};
use crate::utils::Context as _;
pub fn stdin() -> Stdin {
Stdin(Mutex::new(State::Idle(Some(Inner {
stdin: std::io::stdin(),
line: String::new(),
buf: Vec::new(),
last_op: None,
}))))
}
#[derive(Debug)]
pub struct Stdin(Mutex<State>);
#[derive(Debug)]
enum State {
Idle(Option<Inner>),
Busy(JoinHandle<State>),
}
#[derive(Debug)]
struct Inner {
stdin: std::io::Stdin,
line: String,
buf: Vec<u8>,
last_op: Option<Operation>,
}
#[derive(Debug)]
enum Operation {
ReadLine(io::Result<usize>),
Read(io::Result<usize>),
}
impl Stdin {
pub async fn read_line(&self, buf: &mut String) -> io::Result<usize> {
future::poll_fn(|cx| {
let state = &mut *self.0.lock().unwrap();
loop {
match state {
State::Idle(opt) => {
let inner = opt.as_mut().unwrap();
if let Some(Operation::ReadLine(res)) = inner.last_op.take() {
let n = res?;
buf.push_str(&inner.line);
return Poll::Ready(Ok(n));
} else {
let mut inner = opt.take().unwrap();
*state = State::Busy(spawn_blocking(move || {
inner.line.clear();
let res = inner.stdin.read_line(&mut inner.line);
inner.last_op = Some(Operation::ReadLine(res));
State::Idle(Some(inner))
}));
}
}
State::Busy(task) => *state = futures_core::ready!(Pin::new(task).poll(cx)),
}
}
})
.await
.context(|| String::from("could not read line on stdin"))
}
}
impl Read for Stdin {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
let mut state_guard = self.0.lock().unwrap();
let state = &mut *state_guard;
loop {
match state {
State::Idle(opt) => {
let inner = opt.as_mut().unwrap();
if let Some(Operation::Read(res)) = inner.last_op.take() {
let n = res?;
if n <= buf.len() {
buf[..n].copy_from_slice(&inner.buf[..n]);
return Poll::Ready(Ok(n));
}
} else {
let mut inner = opt.take().unwrap();
if inner.buf.len() < buf.len() {
inner.buf.reserve(buf.len() - inner.buf.len());
}
unsafe {
inner.buf.set_len(buf.len());
}
*state = State::Busy(spawn_blocking(move || {
let res = std::io::Read::read(&mut inner.stdin, &mut inner.buf);
inner.last_op = Some(Operation::Read(res));
State::Idle(Some(inner))
}));
}
}
State::Busy(task) => *state = futures_core::ready!(Pin::new(task).poll(cx)),
}
}
}
}
cfg_unix! {
use crate::os::unix::io::{AsRawFd, RawFd};
impl AsRawFd for Stdin {
fn as_raw_fd(&self) -> RawFd {
std::io::stdin().as_raw_fd()
}
}
}
cfg_windows! {
use crate::os::windows::io::{AsRawHandle, RawHandle};
impl AsRawHandle for Stdin {
fn as_raw_handle(&self) -> RawHandle {
std::io::stdin().as_raw_handle()
}
}
}