Write exit code to output (#345)

evan-schott · sjudson · commit f38589ba432a · 2025-02-12T12:37:31.000-05:00
* impl

* indexing bug

* sam + duc comments

* fix
diff --git a/runtime/src/io.rs b/runtime/src/io.rs
@@ -5,7 +5,8 @@ extern crate alloc;
 mod riscv32 {
     extern crate alloc;
     use crate::{
-        ecall, rin, wou, SYS_CYCLE_COUNT, SYS_EXIT, SYS_LOG, SYS_READ_PRIVATE_INPUT, WORD_SIZE,
+        ecall, read_input, write_output, SYS_CYCLE_COUNT, SYS_EXIT, SYS_LOG,
+        SYS_READ_PRIVATE_INPUT, WORD_SIZE,
     };
     use serde::{de::DeserializeOwned, Serialize};
 
@@ -24,6 +25,9 @@ mod riscv32 {
 
     /// Exit the program with the given exit code.
     pub fn exit(exit_code: i32) -> ! {
+        // Write the exit code to the output.
+        let _ = write_output!(0, exit_code);
+        // Exit the program.
         let _ = ecall!(SYS_EXIT, exit_code);
         // Ecall will trigger exit syscall, so we will never return.
         unsafe {
@@ -54,13 +58,13 @@ mod riscv32 {
     pub fn read_public_input<T: DeserializeOwned>() -> Result<T, postcard::Error> {
         // The first word stores the length of the input (in bytes).
         // This length does not take into account the first word itself.
-        let len = rin!(0) as usize;
+        let len = read_input!(0) as usize;
         let padded_len = (len + 3) & !3;
         let mut input = alloc::vec![0u8; padded_len];
 
         // Read the input into the vector.
         for i in 0..((padded_len) / WORD_SIZE) {
-            let word = rin!((i + 1) * WORD_SIZE);
+            let word = read_input!((i + 1) * WORD_SIZE);
             input[i * WORD_SIZE..(i + 1) * WORD_SIZE].copy_from_slice(&word.to_le_bytes());
         }
 
@@ -78,7 +82,7 @@ mod riscv32 {
         // Write bytes in word chunks to output memory.
         bytes.chunks(WORD_SIZE).enumerate().for_each(|(i, chunk)| {
             let word = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
-            wou!(i * WORD_SIZE, word);
+            write_output!((i + 1) * WORD_SIZE, word); // word 0 is reserved for the exit code
         });
 
         Ok(())
diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs
@@ -87,7 +87,7 @@ macro_rules! ecall {
 
 /// Reads word from specified byte address of the public input.
 #[macro_export]
-macro_rules! rin {
+macro_rules! read_input {
     ($i:expr) => {{
         let mut out: u32;
         unsafe {
@@ -106,7 +106,7 @@ macro_rules! rin {
 
 /// Writes word to the public output at specified byte address.
 #[macro_export]
-macro_rules! wou {
+macro_rules! write_output {
     ($i:expr, $data:expr) => {
         unsafe {
             core::arch::asm!(
diff --git a/runtime/src/runtime.rs b/runtime/src/runtime.rs
@@ -1,13 +1,16 @@
 // Nexus VM runtime environment
 // Note: adapted from riscv-rt, which was adapted from cortex-m.
 use crate::alloc::sys_alloc_aligned;
-use crate::{ecall, EXIT_PANIC, EXIT_SUCCESS, SYS_EXIT, SYS_LOG, SYS_OVERWRITE_SP};
+use crate::{ecall, write_output, EXIT_PANIC, EXIT_SUCCESS, SYS_EXIT, SYS_LOG, SYS_OVERWRITE_SP};
 use core::alloc::{GlobalAlloc, Layout};
 use core::panic::PanicInfo;
 
 #[inline(never)]
 #[panic_handler]
 fn panic(_info: &PanicInfo) -> ! {
+    // Write the exit code to the output.
+    let _ = write_output!(0, EXIT_PANIC);
+    // Finish with exit syscall.
     let _ = ecall!(SYS_EXIT, EXIT_PANIC);
     // Ecall will trigger exit syscall, so we will never return.
     unsafe {
@@ -62,6 +65,8 @@ pub unsafe extern "C" fn start_rust() -> u32 {
     // Run the program.
     main();
 
+    // Write the exit code to the output.
+    let _ = write_output!(0, EXIT_SUCCESS);
     // Finish with exit syscall.
     ecall!(SYS_EXIT, EXIT_SUCCESS)
 }
diff --git a/tests/testing-framework/src/lib.rs b/tests/testing-framework/src/lib.rs
@@ -1,5 +1,6 @@
 #[cfg(test)]
 mod test {
+    use nexus_common::constants::WORD_SIZE;
     use nexus_common::cpu::InstructionResult;
 
     use nexus_vm::elf::ElfFile;
@@ -72,6 +73,23 @@ mod test {
         }
     }
 
+    /// Parse the output bytes as exit code and output.
+    fn parse_output<T: DeserializeOwned>(
+        output: Vec<u8>,
+    ) -> Result<(u32, Option<T>), postcard::Error> {
+        // The first 4 bytes store the exit code.
+        let exit_code =
+            u32::from_le_bytes(output[0..4].try_into().expect("Failed to parse exit code"));
+
+        if output.len() == WORD_SIZE {
+            Ok((exit_code, None))
+        } else {
+            // Deserialize the rest as the output.
+            let output: T = from_bytes(&output[WORD_SIZE..]).expect("Deserialization failed");
+            Ok((exit_code, Some(output)))
+        }
+    }
+
     /// Create a temporary directory with a new Cargo project that has nexus_rt as a local dependency.
     fn create_tmp_dir() -> TempDir {
         // Create a temporary directory.
@@ -197,6 +215,7 @@ mod test {
         }
 
         let mut deserialized_output: Option<V> = None;
+        let mut exit_code: u32;
         let ad = vec![0u8; 0xbeef as usize]; // placeholder ad until we have use for it
 
         for elf in elfs {
@@ -210,11 +229,12 @@ mod test {
                     assert_eq!(&emulator.execute(), &io_args.expected_result);
 
                     // Deserialize the output.
-                    if io_args.expected_output.is_some() {
-                        let output_bytes = emulator.get_output().unwrap();
-                        deserialized_output =
-                            Some(from_bytes(&output_bytes).expect("Deserialization failed"));
-                    }
+                    (exit_code, deserialized_output) =
+                        parse_output(emulator.get_output().unwrap()).unwrap();
+                    assert_eq!(
+                        Err(nexus_vm::error::VMError::VMExited(exit_code)),
+                        io_args.expected_result
+                    );
 
                     // Run a second pass with a linear emulator constructed from the harvard emulator.
                     if matches!(emulator_type, EmulatorType::TwoPass) {
@@ -230,11 +250,12 @@ mod test {
                         assert_eq!(&linear_emulator.execute(), &io_args.expected_result);
 
                         // Deserialize the output.
-                        if io_args.expected_output.is_some() {
-                            let output_bytes = linear_emulator.get_output().unwrap();
-                            deserialized_output =
-                                Some(from_bytes(&output_bytes).expect("Deserialization failed"));
-                        }
+                        (exit_code, deserialized_output) =
+                            parse_output(linear_emulator.get_output().unwrap()).unwrap();
+                        assert_eq!(
+                            Err(nexus_vm::error::VMError::VMExited(exit_code)),
+                            io_args.expected_result
+                        );
                     }
                 }
                 EmulatorType::Linear(heap_size, stack_size, program_size) => {
@@ -269,11 +290,12 @@ mod test {
                     assert_eq!(&emulator.execute(), &io_args.expected_result);
 
                     // Deserialize the output.
-                    if io_args.expected_output.is_some() {
-                        let output_bytes = emulator.get_output().unwrap();
-                        deserialized_output =
-                            Some(from_bytes(&output_bytes).expect("Deserialization failed"));
-                    }
+                    (exit_code, deserialized_output) =
+                        parse_output(emulator.get_output().unwrap()).unwrap();
+                    assert_eq!(
+                        Err(nexus_vm::error::VMError::VMExited(exit_code)),
+                        io_args.expected_result
+                    );
                 }
             };
         }
diff --git a/vm/src/emulator/executor.rs b/vm/src/emulator/executor.rs
@@ -520,7 +520,7 @@ impl LinearEmulator {
                     * WORD_SIZE) as u32,
                 ad.len() as u32,
                 public_input.len() as u32,
-                output_memory.len() as u32,
+                (output_memory.len() - WORD_SIZE) as u32, // Exclude the first word which is the exit code
             )
             .unwrap();
 
@@ -615,10 +615,13 @@ impl LinearEmulator {
             let _ = memory.add_fixed_ro(&input_memory).unwrap();
         }
 
-        let output_len = (memory_layout.public_output_end() - memory_layout.panic()) as usize; // we include panic in the output segment
+        let output_len = (memory_layout.public_output_end() - memory_layout.exit_code()) as usize; // we include the exit code in the output segment
         if output_len > 0 {
-            let output_memory =
-                FixedMemory::<WO>::from_vec(memory_layout.panic(), output_len, vec![0; output_len]);
+            let output_memory = FixedMemory::<WO>::from_vec(
+                memory_layout.exit_code(),
+                output_len,
+                vec![0; output_len],
+            );
             let _ = memory.add_fixed_wo(&output_memory).unwrap();
         }
 
@@ -656,7 +659,7 @@ impl LinearEmulator {
                 8,
                 &[
                     memory_layout.public_input_start(),
-                    memory_layout.public_output_start(),
+                    memory_layout.exit_code(), // The exit code is the first word of the output
                 ],
             ))
             .unwrap();
@@ -682,7 +685,7 @@ impl LinearEmulator {
     pub fn get_output(&self) -> Result<Vec<u8>, MemoryError> {
         Ok(self.memory.segment_bytes(
             (Modes::WO as usize, 0),
-            self.memory_layout.public_output_start(),
+            self.memory_layout.exit_code(),
             Some(self.memory_layout.public_output_end()),
         )?)
     }
diff --git a/vm/src/emulator/layout.rs b/vm/src/emulator/layout.rs
@@ -8,8 +8,8 @@ use serde::{Deserialize, Serialize};
 pub struct LinearMemoryLayout {
     // start of the public input
     public_input: u32,
-    // location of the panic byte
-    panic: u32,
+    // location of the exit code
+    exit_code: u32,
     // start of the public output
     public_output: u32,
     // start of the associated data hash
@@ -49,8 +49,8 @@ impl LinearMemoryLayout {
             || self.gap_start() < self.heap_start()
             || self.heap_start() < self.ad_start()
             || self.ad_start() < self.public_output_start()
-            || self.public_output_start() <= self.panic()
-            || self.panic() <= self.public_input_start() // First word of input stores input length, so must be non-empty
+            || self.public_output_start() <= self.exit_code()
+            || self.exit_code() <= self.public_input_start() // First word of input stores input length, so must be non-empty
             || self.public_input_start() <= self.program_start() // Program assumed to be non-empty
             || self.program_start() <= self.public_output_start_location()
         {
@@ -69,8 +69,8 @@ impl LinearMemoryLayout {
         ad_size: u32,
     ) -> Self {
         let public_input = ELF_TEXT_START + program_size;
-        let panic = public_input + public_input_size + WORD_SIZE as u32;
-        let public_output = panic + WORD_SIZE as u32;
+        let exit_code = public_input + public_input_size + WORD_SIZE as u32;
+        let public_output = exit_code + WORD_SIZE as u32;
         let ad = public_output + public_output_size;
         let heap = ad + ad_size;
         let gap = heap + max_heap_size;
@@ -79,7 +79,7 @@ impl LinearMemoryLayout {
 
         Self {
             public_input,
-            panic,
+            exit_code,
             public_output,
             ad,
             heap,
@@ -143,11 +143,11 @@ impl LinearMemoryLayout {
     }
 
     pub fn public_input_end(&self) -> u32 {
-        self.panic
+        self.exit_code
     }
 
-    pub fn panic(&self) -> u32 {
-        self.panic
+    pub fn exit_code(&self) -> u32 {
+        self.exit_code
     }
 
     pub fn public_output_start(&self) -> u32 {
diff --git a/vm/src/memory/fixed.rs b/vm/src/memory/fixed.rs
@@ -93,7 +93,7 @@ impl<M: Mode> FixedMemory<M> {
         let s = (start - self.base_address) / WORD_SIZE as u32;
 
         if let Some(mut e) = end {
-            e -= self.base_address;
+            e = (e - self.base_address) / WORD_SIZE as u32;
             &self.vec[s as usize..e as usize]
         } else {
             &self.vec[s as usize..]
