diff --git a/riscv/runtime/src/fmt.rs b/riscv/runtime/src/fmt.rs new file mode 100644 index 000000000..a457668f2 --- /dev/null +++ b/riscv/runtime/src/fmt.rs @@ -0,0 +1,40 @@ +use core::arch::asm; +use core::fmt; + +#[macro_export] +macro_rules! print { + ($($arg:tt)+) => {{ + $crate::fmt::print_args(format_args!( $($arg)+)); + }}; +} + +pub fn print_args(args: fmt::Arguments) { + fmt::write(&mut ProverWriter {}, args).unwrap(); +} + +struct ProverWriter {} + +impl fmt::Write for ProverWriter { + fn write_str(&mut self, s: &str) -> fmt::Result { + print_str(s); + Ok(()) + } +} + +pub fn print_str(s: &str) { + // DEAR DEV, please don't allow this function to panic. + // + // This is called from the panic handler. + for b in s.bytes() { + print_prover_char(b) + } +} + +#[inline] +fn print_prover_char(c: u8) { + let mut value = c as u32; + #[allow(unused_assignments)] + unsafe { + asm!("ebreak", lateout("a0") value, in("a0") value); + } +} diff --git a/riscv/runtime/src/lib.rs b/riscv/runtime/src/lib.rs index 9cccef6af..c6276f798 100644 --- a/riscv/runtime/src/lib.rs +++ b/riscv/runtime/src/lib.rs @@ -1,19 +1,32 @@ #![no_std] -#![feature(start, alloc_error_handler)] +#![feature( + start, + alloc_error_handler, + maybe_uninit_write_slice, + round_char_boundary +)] use core::arch::asm; use core::panic::PanicInfo; +use crate::fmt::print_str; + mod allocator; -mod print; -pub use print::print; +pub mod fmt; #[panic_handler] -fn panic(panic: &PanicInfo<'_>) -> ! { - print(format_args!("{panic}")); - unsafe { - asm!("unimp"); +unsafe fn panic(panic: &PanicInfo<'_>) -> ! { + static mut IS_PANICKING: bool = false; + + if !IS_PANICKING { + IS_PANICKING = true; + + print!("{panic}\n"); + } else { + print_str("Panic handler has panicked! Things are very dire indeed...\n"); } + + asm!("unimp"); loop {} } diff --git a/riscv/runtime/src/print.rs b/riscv/runtime/src/print.rs deleted file mode 100644 index ce8488821..000000000 --- a/riscv/runtime/src/print.rs +++ /dev/null @@ -1,81 +0,0 @@ -use core::arch::asm; -use core::fmt; -use core::mem::MaybeUninit; - -// #[macro_export] -// macro_rules! print { -// ($($arg:tt)+) => (print_args(format_args!( $($arg)+))) -// } - -// TODO turn this into a macro -pub fn print(args: fmt::Arguments) { - const BUF_SIZE: usize = 1024; - let mut buf = unsafe { MaybeUninit::<[MaybeUninit; BUF_SIZE]>::uninit().assume_init() }; - let _s: &str = buf_formatter::format(&mut buf, args).unwrap(); - print_prover(_s); -} - -#[inline] -fn print_prover(s: &str) { - for b in s.bytes() { - print_prover_char(b) - } -} - -#[inline] -fn print_prover_char(c: u8) { - let mut value = c as u32; - #[allow(unused_assignments)] - unsafe { - asm!("ebreak", lateout("a0") value, in("a0") value); - } -} - -mod buf_formatter { - use core::fmt; - use core::mem::MaybeUninit; - - pub struct BufFormatter<'a> { - buffer: &'a mut [MaybeUninit], - used: usize, - } - - impl<'a> BufFormatter<'a> { - pub fn new(buffer: &'a mut [MaybeUninit]) -> Self { - Self { buffer, used: 0 } - } - - pub fn as_str(&self) -> &'a str { - unsafe { - // This is safe because everything until self.used has been initialized. - let buf = - &*(&self.buffer[..self.used] as *const [MaybeUninit] as *const [u8]); - // we only concatenate str, so the result must be valid utf8 as well. - core::str::from_utf8_unchecked(buf) - } - } - } - - impl<'a> fmt::Write for BufFormatter<'a> { - fn write_str(&mut self, s: &str) -> fmt::Result { - let raw_s = s.as_bytes(); - let write_len = raw_s.len(); - if write_len > self.buffer.len() - self.used { - return Err(fmt::Error); - } - let s_uninit: &[MaybeUninit] = unsafe { core::mem::transmute(&raw_s[..write_len]) }; - self.buffer[self.used..][..write_len].copy_from_slice(s_uninit); - self.used += write_len; - Ok(()) - } - } - - pub fn format<'a>( - buffer: &'a mut [MaybeUninit], - args: fmt::Arguments, - ) -> Result<&'a str, fmt::Error> { - let mut w = BufFormatter::new(buffer); - fmt::write(&mut w, args)?; - Ok(w.as_str()) - } -} diff --git a/riscv/src/lib.rs b/riscv/src/lib.rs index a82614336..7b84f5f0f 100644 --- a/riscv/src/lib.rs +++ b/riscv/src/lib.rs @@ -164,9 +164,9 @@ runtime = {{ path = "./runtime" }} ) .unwrap(); - let mut print_file = runtime_file.clone(); - print_file.push("print.rs"); - fs::write(print_file, include_bytes!("../runtime/src/print.rs")).unwrap(); + let mut fmt_file = runtime_file.clone(); + fmt_file.push("fmt.rs"); + fs::write(fmt_file, include_bytes!("../runtime/src/fmt.rs")).unwrap(); compile_rust_crate_to_riscv_asm(cargo_file.to_str().unwrap()) } diff --git a/riscv/tests/riscv_data/vec_median/src/lib.rs b/riscv/tests/riscv_data/vec_median/src/lib.rs index 1165cfa12..4ef5494f1 100644 --- a/riscv/tests/riscv_data/vec_median/src/lib.rs +++ b/riscv/tests/riscv_data/vec_median/src/lib.rs @@ -16,7 +16,7 @@ extern crate alloc; use alloc::vec::Vec; -use runtime::get_prover_input; +use runtime::{get_prover_input, print}; #[no_mangle] fn main() { @@ -33,5 +33,6 @@ fn main() { (vec[half - 1] + vec[half]) / 2 }; + print!("Found median of {median}\n"); assert_eq!(median, expected); }