tg_console/
lib.rs

1//! 提供可定制实现的 `print!`、`println!` 和 `log::Log`。
2
3#![no_std]
4#![deny(warnings, missing_docs)]
5
6use core::{
7    fmt::{self, Write},
8    str::FromStr,
9};
10use spin::Once;
11
12/// 向用户提供 `log`。
13pub extern crate log;
14
15/// 这个接口定义了向控制台“输出”这件事。
16pub trait Console: Sync {
17    /// 向控制台放置一个字符。
18    fn put_char(&self, c: u8);
19
20    /// 向控制台放置一个字符串。
21    ///
22    /// 如果使用了锁，覆盖这个实现以免反复获取和释放锁。
23    #[inline]
24    fn put_str(&self, s: &str) {
25        for c in s.bytes() {
26            self.put_char(c);
27        }
28    }
29}
30
31/// 库找到输出的方法：保存一个对象引用，这是一种单例。
32static CONSOLE: Once<&'static dyn Console> = Once::new();
33
34/// 打印缓冲区大小。
35const BUFFER_SIZE: usize = 64;
36
37/// 打印缓冲区，用于收集格式化输出后一次性输出，避免抢占导致输出交错。
38struct PrintBuffer {
39    buffer: [u8; BUFFER_SIZE],
40    pos: usize,
41}
42
43impl PrintBuffer {
44    /// 创建一个新的打印缓冲区。
45    const fn new() -> Self {
46        Self {
47            buffer: [0u8; BUFFER_SIZE],
48            pos: 0,
49        }
50    }
51
52    /// 刷新缓冲区，将内容输出到控制台。
53    fn flush(&mut self) {
54        if self.pos > 0 {
55            if let Some(console) = CONSOLE.get() {
56                // SAFETY: buffer 中的内容是从 str 写入的，保证是有效的 UTF-8
57                let s = unsafe { core::str::from_utf8_unchecked(&self.buffer[..self.pos]) };
58                console.put_str(s);
59            }
60            self.pos = 0;
61        }
62    }
63
64    /// 写入字符串到缓冲区，必要时刷新。
65    fn write(&mut self, s: &str) {
66        let bytes = s.as_bytes();
67        let mut offset = 0;
68
69        while offset < bytes.len() {
70            let remaining_buffer = BUFFER_SIZE - self.pos;
71            let remaining_input = bytes.len() - offset;
72            let to_copy = remaining_buffer.min(remaining_input);
73
74            self.buffer[self.pos..self.pos + to_copy]
75                .copy_from_slice(&bytes[offset..offset + to_copy]);
76            self.pos += to_copy;
77            offset += to_copy;
78
79            // 缓冲区满了，刷新
80            if self.pos == BUFFER_SIZE {
81                self.flush();
82            }
83        }
84    }
85}
86
87impl Write for PrintBuffer {
88    fn write_str(&mut self, s: &str) -> fmt::Result {
89        self.write(s);
90        Ok(())
91    }
92}
93
94/// 用户调用这个函数设置输出的方法。
95pub fn init_console(console: &'static dyn Console) {
96    CONSOLE.call_once(|| console);
97    log::set_logger(&Logger).unwrap();
98}
99
100/// 根据环境变量设置日志级别。
101pub fn set_log_level(env: Option<&str>) {
102    use log::LevelFilter as Lv;
103    log::set_max_level(env.and_then(|s| Lv::from_str(s).ok()).unwrap_or(Lv::Trace));
104}
105
106/// 打印一些测试信息。
107pub fn test_log() {
108    println!(
109        r"
110   ______                       __
111  / ____/___  ____  _________  / /__
112 / /   / __ \/ __ \/ ___/ __ \/ / _ \
113/ /___/ /_/ / / / (__  ) /_/ / /  __/
114\____/\____/_/ /_/____/\____/_/\___/
115===================================="
116    );
117    log::trace!("LOG TEST >> Hello, world!");
118    log::debug!("LOG TEST >> Hello, world!");
119    log::info!("LOG TEST >> Hello, world!");
120    log::warn!("LOG TEST >> Hello, world!");
121    log::error!("LOG TEST >> Hello, world!");
122    println!();
123}
124
125/// 打印。
126///
127/// 给宏用的，用户不会直接调它。
128/// 使用栈上缓冲区收集格式化输出，然后一次性输出到控制台。
129#[doc(hidden)]
130#[inline]
131pub fn _print(args: fmt::Arguments) {
132    let mut buffer = PrintBuffer::new();
133    buffer.write_fmt(args).unwrap();
134    buffer.flush();
135}
136
137/// 格式化打印。
138#[macro_export]
139macro_rules! print {
140    ($($arg:tt)*) => {
141        $crate::_print(core::format_args!($($arg)*));
142    }
143}
144
145/// 格式化打印并换行。
146#[macro_export]
147macro_rules! println {
148    () => ($crate::print!("\n"));
149    ($($arg:tt)*) => {{
150        $crate::_print(core::format_args!($($arg)*));
151        $crate::println!();
152    }}
153}
154
155/// 这个 Unit struct 是 `core::fmt` 要求的。
156struct Logger;
157
158/// 实现 [`Write`] trait，格式化的基础。
159impl Write for Logger {
160    #[inline]
161    fn write_str(&mut self, s: &str) -> Result<(), fmt::Error> {
162        let _ = CONSOLE.get().unwrap().put_str(s);
163        Ok(())
164    }
165}
166
167/// 实现 `log::Log` trait，提供分级日志。
168impl log::Log for Logger {
169    #[inline]
170    fn enabled(&self, _metadata: &log::Metadata) -> bool {
171        true
172    }
173
174    #[inline]
175    fn log(&self, record: &log::Record) {
176        use log::Level::*;
177        let color_code: u8 = match record.level() {
178            Error => 31,
179            Warn => 93,
180            Info => 34,
181            Debug => 32,
182            Trace => 90,
183        };
184        println!(
185            "\x1b[{color_code}m[{:>5}] {}\x1b[0m",
186            record.level(),
187            record.args(),
188        );
189    }
190
191    fn flush(&self) {}
192}
tg_console/lib.rs

tg_console/
lib.rs
