ax_std/io/
stdio.rs

1#[cfg(feature = "alloc")]
2use alloc::{string::String, vec::Vec};
3
4use ax_lazyinit::LazyInit;
5
6use crate::{
7    io::{self, BufReader, prelude::*},
8    sync::{Mutex, MutexGuard},
9};
10
11struct StdinRaw;
12struct StdoutRaw;
13
14impl Read for StdinRaw {
15    // Non-blocking read, returns number of bytes read.
16    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
17        let mut read_len = 0;
18        while read_len < buf.len() {
19            let len = ax_api::stdio::ax_console_read_bytes(buf[read_len..].as_mut())?;
20            if len == 0 {
21                break;
22            }
23            read_len += len;
24        }
25        Ok(read_len)
26    }
27}
28
29impl Write for StdoutRaw {
30    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
31        Ok(ax_api::stdio::ax_console_write_bytes(buf)?)
32    }
33    fn flush(&mut self) -> io::Result<()> {
34        ax_api::stdio::ax_console_flush()?;
35        Ok(())
36    }
37}
38
39/// A handle to the standard input stream of a process.
40pub struct Stdin {
41    inner: &'static Mutex<BufReader<StdinRaw>>,
42}
43
44/// A locked reference to the [`Stdin`] handle.
45pub struct StdinLock<'a> {
46    inner: MutexGuard<'a, BufReader<StdinRaw>>,
47}
48
49impl Stdin {
50    /// Locks this handle to the standard input stream, returning a readable
51    /// guard.
52    ///
53    /// The lock is released when the returned lock goes out of scope. The
54    /// returned guard also implements the [`Read`] and [`BufRead`] traits for
55    /// accessing the underlying data.
56    pub fn lock(&self) -> StdinLock<'static> {
57        // Locks this handle with 'static lifetime. This depends on the
58        // implementation detail that the underlying `Mutex` is static.
59        StdinLock {
60            inner: self.inner.lock(),
61        }
62    }
63
64    /// Locks this handle and reads a line of input, appending it to the specified buffer.
65    #[cfg(feature = "alloc")]
66    pub fn read_line(&self, buf: &mut String) -> io::Result<usize> {
67        self.inner.lock().read_line(buf)
68    }
69}
70
71impl Read for Stdin {
72    // Block until at least one byte is read.
73    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
74        let read_len = self.inner.lock().read(buf)?;
75        if buf.is_empty() || read_len > 0 {
76            return Ok(read_len);
77        }
78        // Sleep until the runtime RX worker publishes progress, then retry.
79        loop {
80            ax_api::stdio::ax_console_wait_readable()?;
81            let read_len = self.inner.lock().read(buf)?;
82            if read_len > 0 {
83                return Ok(read_len);
84            }
85        }
86    }
87}
88
89impl Read for StdinLock<'_> {
90    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
91        self.inner.read(buf)
92    }
93}
94
95impl BufRead for StdinLock<'_> {
96    fn fill_buf(&mut self) -> io::Result<&[u8]> {
97        self.inner.fill_buf()
98    }
99
100    fn consume(&mut self, n: usize) {
101        self.inner.consume(n)
102    }
103
104    #[cfg(feature = "alloc")]
105    fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> io::Result<usize> {
106        self.inner.read_until(byte, buf)
107    }
108
109    #[cfg(feature = "alloc")]
110    fn read_line(&mut self, buf: &mut String) -> io::Result<usize> {
111        self.inner.read_line(buf)
112    }
113}
114
115/// A handle to the global standard output stream of the current process.
116pub struct Stdout {
117    inner: &'static Mutex<StdoutRaw>,
118}
119
120/// A locked reference to the [`Stdout`] handle.
121pub struct StdoutLock<'a> {
122    inner: MutexGuard<'a, StdoutRaw>,
123}
124
125impl Stdout {
126    /// Locks this handle to the standard output stream, returning a writable
127    /// guard.
128    ///
129    /// The lock is released when the returned lock goes out of scope. The
130    /// returned guard also implements the `Write` trait for writing data.
131    pub fn lock(&self) -> StdoutLock<'static> {
132        StdoutLock {
133            inner: self.inner.lock(),
134        }
135    }
136}
137
138impl Write for Stdout {
139    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
140        self.inner.lock().write(buf)
141    }
142    fn flush(&mut self) -> io::Result<()> {
143        self.inner.lock().flush()
144    }
145}
146
147impl Write for StdoutLock<'_> {
148    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
149        self.inner.write(buf)
150    }
151    fn flush(&mut self) -> io::Result<()> {
152        self.inner.flush()
153    }
154}
155
156/// Constructs a new handle to the standard input of the current process.
157pub fn stdin() -> Stdin {
158    static INSTANCE: LazyInit<Mutex<BufReader<StdinRaw>>> = LazyInit::new();
159    if !INSTANCE.is_inited() {
160        INSTANCE.init_once(Mutex::new(BufReader::new(StdinRaw)));
161    }
162    Stdin { inner: &INSTANCE }
163}
164
165/// Constructs a new handle to the standard output of the current process.
166pub fn stdout() -> Stdout {
167    static INSTANCE: LazyInit<Mutex<StdoutRaw>> = LazyInit::new();
168    if !INSTANCE.is_inited() {
169        INSTANCE.init_once(Mutex::new(StdoutRaw));
170    }
171    Stdout { inner: &INSTANCE }
172}
173
174#[doc(hidden)]
175pub fn __print_impl(args: core::fmt::Arguments) {
176    if cfg!(feature = "smp") {
177        // The runtime serializes formatted user output on the sleepable TTY
178        // path; kernel logs use a separate non-blocking mailbox.
179        ax_api::stdio::ax_console_write_fmt(args).unwrap();
180    } else {
181        stdout().lock().write_fmt(args).unwrap();
182    }
183}
ax_std/io/stdio.rs

ax_std/io/
stdio.rs
