Skip to main content

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        ax_api::stdio::ax_console_write_bytes(buf)
32    }
33    fn flush(&mut self) -> io::Result<()> {
34        Ok(())
35    }
36}
37
38/// A handle to the standard input stream of a process.
39pub struct Stdin {
40    inner: &'static Mutex<BufReader<StdinRaw>>,
41}
42
43/// A locked reference to the [`Stdin`] handle.
44pub struct StdinLock<'a> {
45    inner: MutexGuard<'a, BufReader<StdinRaw>>,
46}
47
48impl Stdin {
49    /// Locks this handle to the standard input stream, returning a readable
50    /// guard.
51    ///
52    /// The lock is released when the returned lock goes out of scope. The
53    /// returned guard also implements the [`Read`] and [`BufRead`] traits for
54    /// accessing the underlying data.
55    pub fn lock(&self) -> StdinLock<'static> {
56        // Locks this handle with 'static lifetime. This depends on the
57        // implementation detail that the underlying `Mutex` is static.
58        StdinLock {
59            inner: self.inner.lock(),
60        }
61    }
62
63    /// Locks this handle and reads a line of input, appending it to the specified buffer.
64    #[cfg(feature = "alloc")]
65    pub fn read_line(&self, buf: &mut String) -> io::Result<usize> {
66        self.inner.lock().read_line(buf)
67    }
68}
69
70impl Read for Stdin {
71    // Block until at least one byte is read.
72    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
73        let read_len = self.inner.lock().read(buf)?;
74        if buf.is_empty() || read_len > 0 {
75            return Ok(read_len);
76        }
77        // try again until we got something
78        loop {
79            let read_len = self.inner.lock().read(buf)?;
80            if read_len > 0 {
81                return Ok(read_len);
82            }
83            crate::thread::yield_now();
84        }
85    }
86}
87
88impl Read for StdinLock<'_> {
89    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
90        self.inner.read(buf)
91    }
92}
93
94impl BufRead for StdinLock<'_> {
95    fn fill_buf(&mut self) -> io::Result<&[u8]> {
96        self.inner.fill_buf()
97    }
98
99    fn consume(&mut self, n: usize) {
100        self.inner.consume(n)
101    }
102
103    #[cfg(feature = "alloc")]
104    fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> io::Result<usize> {
105        self.inner.read_until(byte, buf)
106    }
107
108    #[cfg(feature = "alloc")]
109    fn read_line(&mut self, buf: &mut String) -> io::Result<usize> {
110        self.inner.read_line(buf)
111    }
112}
113
114/// A handle to the global standard output stream of the current process.
115pub struct Stdout {
116    inner: &'static Mutex<StdoutRaw>,
117}
118
119/// A locked reference to the [`Stdout`] handle.
120pub struct StdoutLock<'a> {
121    inner: MutexGuard<'a, StdoutRaw>,
122}
123
124impl Stdout {
125    /// Locks this handle to the standard output stream, returning a writable
126    /// guard.
127    ///
128    /// The lock is released when the returned lock goes out of scope. The
129    /// returned guard also implements the `Write` trait for writing data.
130    pub fn lock(&self) -> StdoutLock<'static> {
131        StdoutLock {
132            inner: self.inner.lock(),
133        }
134    }
135}
136
137impl Write for Stdout {
138    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
139        self.inner.lock().write(buf)
140    }
141    fn flush(&mut self) -> io::Result<()> {
142        self.inner.lock().flush()
143    }
144}
145
146impl Write for StdoutLock<'_> {
147    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
148        self.inner.write(buf)
149    }
150    fn flush(&mut self) -> io::Result<()> {
151        self.inner.flush()
152    }
153}
154
155/// Constructs a new handle to the standard input of the current process.
156pub fn stdin() -> Stdin {
157    static INSTANCE: LazyInit<Mutex<BufReader<StdinRaw>>> = LazyInit::new();
158    if !INSTANCE.is_inited() {
159        INSTANCE.init_once(Mutex::new(BufReader::new(StdinRaw)));
160    }
161    Stdin { inner: &INSTANCE }
162}
163
164/// Constructs a new handle to the standard output of the current process.
165pub fn stdout() -> Stdout {
166    static INSTANCE: LazyInit<Mutex<StdoutRaw>> = LazyInit::new();
167    if !INSTANCE.is_inited() {
168        INSTANCE.init_once(Mutex::new(StdoutRaw));
169    }
170    Stdout { inner: &INSTANCE }
171}
172
173#[doc(hidden)]
174pub fn __print_impl(args: core::fmt::Arguments) {
175    if cfg!(feature = "smp") {
176        // synchronize using the lock in ax-log, to avoid interleaving
177        // with kernel logs
178        ax_api::stdio::ax_console_write_fmt(args).unwrap();
179    } else {
180        stdout().lock().write_fmt(args).unwrap();
181    }
182}