【Rust】在no_std里进行一个print

tangptr@126.com

前言

Rust的print!和println!在#![no_std]里是不可用的，得自己去实现。

格式化参数

在Rust中，函数是不可以有可变数量的参数的，但是宏可以！Rust提供了名为format_args的宏。它可以将输入进来的参数构造为一个Arguments结构体。与这个结构体直接相关的，就是Write trait。你需要实现这个trait里的write_str方法来把Arguments变成字符串。
对应到C语言，可以粗略地理解为Arguments是va_list，format_args是va_start，而Write::write_str是vsnprintf中针对不同类型的打印函数（你会发现Rust这个套路比C的vsnprintf好使多了）。
接下来实现print和println宏：

macro_rules! print
{
    ($($arg:tt)*) =>
    {
        (internal_print(format_args!($($arg)*)))        
    };
}

macro_rules! println
{
    () =>
    {
        print!("\n")
    };
    ($($arg:tt)*) =>
    {
        print!("{}\n",format_args!($($arg)*))
    };
}

而internal_print函数的实现方法稍后再说。

栈上Print

顾名思义，就是把Arguments写进栈上的字符串。这个方法会限制一次print的长度，但是可以避免内存分配，性能较高，而且在Windows驱动里的高IRQL的条件下非常实用。这里采用512是也是因为Windows内核的DbgPrint一次最多只能发送512个字节。微软的windows-drivers-rs里的栈上print就是我实现的。
首先我们定义一个栈上缓冲区：

pub struct FormatBuffer
{
    used:usize,
    buffer:MaybeUninit<[u8;512]>
}

这里用MaybeUninit可以避免无意义的初始化行为（直接[0;N]的开销很大）。接下来实现一下Default trait用来初始化：

impl Default for FormatBuffer
{
    fn default() -> Self
    {
        Self
        {
            used:0,
            buffer:MaybeUninit::uninit()
        }
    }
}

接下来实现Write trait：

impl fmt::Write for FormatBuffer
{
    fn write_str(&mut self, s: &str) -> fmt::Result
    {
        let remainder=unsafe{&mut self.buffer.assume_init_mut()[self.used..]};
        let current=s.as_bytes();
        if remainder.len()<current.len()
        {
            return Err(fmt::Error);
        }
        remainder[..current.len()].copy_from_slice(current);
        self.used+=current.len();
        Ok(())
    }
}

顺便实现一下as_str方法，这样一来打印起来更直观：

impl FormatBuffer
{
    pub fn as_slice(&self)->&[u8]
    {
        unsafe
        {
            &self.buffer.assume_init_ref()[..self.used]
        }
    }

    pub fn as_str(&self)->&str
    {
        unsafe
        {
            str::from_utf8_unchecked(self.as_slice())
        }
    }
}

以Windows平台为例，实现方法如下：

static mut STDOUT_HANDLE: LazyCell<HANDLE>=LazyCell::new(|| unsafe{GetStdHandle(STD_OUTPUT_HANDLE)});

pub fn internal_print(args:fmt::Arguments)
{
    let mut w=FormatBuffer::default();
    let r=fmt::write(&mut w,args);
    if r.is_ok()
    {
        let s=w.as_str();
        let mut utf16_buff:MaybeUninit<[u16;512]>=MaybeUninit::uninit();
        let mut written_length:u32=0;
        let mut i:usize=0;
        for c in s.encode_utf16()
        {
            unsafe
            {
                utf16_buff.assume_init_mut()[i]=c;
            }
            i+=1;
        }
        unsafe
        {
            WriteConsoleW(*STDOUT_HANDLE,utf16_buff.assume_init_ref().as_ptr(),i as u32,&raw mut written_length,null());
        }
    }
}

注意，Rust的字符串是UTF-8的，不能直接传给WriteConsoleA函数，否则打印汉字什么的就操蛋了。需要先转为UTF-16然后传给WriteConsoleW。
这里用了LazyCell类型来存放stdout的句柄，可以降低初始化全局变量的复杂度。但注意，这里为了演示，用了LazyCell（多线程不安全，所以只能用static mut声明它）。涉及多线程开发时需要用LazyLock（很不幸，no_std环境里没有，需要自己实现锁），或者在程序启动时提前用force方法对其提前初始化。

堆上Print

顾名思义，就是把Arguments写进堆上的字符串，这种方法不限制一次print的最大长度。我们可以直接用alloc库里提供的String类型，因为它已经实现过Write trait了。我们实现一下internal_print就可以了：

pub fn internal_print(args:fmt::Arguments)
{
    let mut w=String::new();
    let r=fmt::write(&mut w,args);
    if r.is_ok()
    {
        let v:Vec<u16>=w.encode_utf16().collect();
        let mut written_length:u32=0;
        unsafe
        {
            WriteConsoleW(*STDOUT_HANDLE,v.as_ptr(),v.len() as u32,&raw mut written_length,null());
        }
    }
}

但是在no_std的环境下，需要实现一个内存分配器才能用堆上分配。因此需要实现GlobalAlloc trait。以HeapAlloc为例：

static mut PROCESS_HEAP: LazyCell<HANDLE>=LazyCell::new(|| unsafe{GetProcessHeap()});
struct WindowsAllocator;

unsafe impl GlobalAlloc for WindowsAllocator
{
    unsafe fn alloc(&self,layout: Layout)->*mut u8
    {
        unsafe
        {
            HeapAlloc(*PROCESS_HEAP,0,layout.size()).cast()
        }
    }

    unsafe fn dealloc(&self,ptr:*mut u8,_layout: Layout)
    {
        unsafe
        {
            HeapFree(*PROCESS_HEAP,0,ptr.cast());
        }
    }

    unsafe fn realloc(&self,ptr:*mut u8,_layout: Layout,new_size:usize)->*mut u8
    {
        unsafe
        {
            HeapReAlloc(*PROCESS_HEAP,0,ptr.cast(),new_size).cast()
        }
    }

    unsafe fn alloc_zeroed(&self,layout: Layout)->*mut u8
    {
        unsafe
        {
            HeapAlloc(*PROCESS_HEAP,HEAP_ZERO_MEMORY,layout.size()).cast()
        }
    }
}

其中，realloc和alloc_zeroed是可选的。如果你用的内存分配器没有专门的优化，就不需要实现。（比如realloc的默认实现是分配一个大的buffer，copy旧的过去之后，再free掉。）
注意，以上实现不考虑内存对齐的问题！但是对于字符串来说，对齐粒度无所谓。以后再单开一篇讲内存分配。

结语

你会发现，其实也就是internal_print函数需要根据平台不同而调用不同的API（以及堆上print需要的分配器），其他的都是通用的。