From 2f13506971b33b9102d648c14f9503d1bdad2f55 Mon Sep 17 00:00:00 2001
From: Clo91eaf <Clo91eaf@qq.com>
Date: Wed, 24 Jul 2024 14:49:04 +0800
Subject: [PATCH] [rocketemu] exit with right EXIT_CODE check

---
 rocketemu/driver/src/dpi.rs | 12 +++++++-----
 rocketemu/driver/src/sim.rs | 19 ++++++++++---------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/rocketemu/driver/src/dpi.rs b/rocketemu/driver/src/dpi.rs
index 945b545be..8103ebb1c 100644
--- a/rocketemu/driver/src/dpi.rs
+++ b/rocketemu/driver/src/dpi.rs
@@ -22,17 +22,18 @@ pub type SvBitVecVal = u32;
 unsafe fn load_from_payload(
   payload: &*const SvBitVecVal,
   aw_size: c_longlong,
-  data_width: u32,
+  data_width: usize,
+  dlen: usize,
 ) -> (Vec<bool>, &[u8]) {
   let src = *payload as *mut u8;
-  let data_width_in_byte = (data_width / 8) as usize;
-  let strb_width_in_byte = data_width_in_byte.div_ceil(8); // ceil divide by 8 to get byte width
+  let data_width_in_byte = dlen / 8;
+  let strb_width_in_byte = dlen / data_width;
   let payload_size_in_byte = strb_width_in_byte + data_width_in_byte; // data width in byte
   let byte_vec = std::slice::from_raw_parts(src, payload_size_in_byte);
   let strobe = &byte_vec[0..strb_width_in_byte];
   let data = &byte_vec[strb_width_in_byte..];
 
-  let strb_width_in_bit = std::cmp::min(8, data_width_in_byte);
+  let strb_width_in_bit = data_width / 8;
   let masks: Vec<bool> = strobe
     .into_iter()
     .flat_map(|strb| {
@@ -97,7 +98,8 @@ unsafe extern "C" fn axi_write_loadStoreAXI_rs(
   );
 
   let sim = &mut *(target as *mut Simulator);
-  let (strobe, data) = load_from_payload(&payload, 1 << awsize, sim.dlen);
+  let data_width = 32; // TODO: get from sim
+  let (strobe, data) = load_from_payload(&payload, 1 << awsize, data_width, sim.dlen as usize);
   sim.axi_write(awaddr as u32, &strobe, data);
 }
 
diff --git a/rocketemu/driver/src/sim.rs b/rocketemu/driver/src/sim.rs
index bd5b2f3c5..2c9eb45f5 100644
--- a/rocketemu/driver/src/sim.rs
+++ b/rocketemu/driver/src/sim.rs
@@ -25,6 +25,7 @@ pub(crate) struct AxiReadPayload {
 }
 
 const EXIT_POS: u32 = 0x4000_0000;
+const EXIT_CODE: u32 = 0xdead_beef;
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
@@ -263,15 +264,15 @@ impl Simulator {
       get_t()
     );
     let data_hex = hex::encode(data);
-    info!(
-      "[{}] axi_write (addr={addr:#x}, data={data_hex})",
-      get_t()
-    );
-
-    if addr == EXIT_POS {
-      info!("exit with code: {:x?}", data);
-      quit();
-      return;
+    info!("[{}] axi_write (addr={addr:#x}, data={data_hex})", get_t());
+
+    if addr == EXIT_POS && data.len() >= 4 {
+      let exit_code = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
+      if exit_code == EXIT_CODE {
+        info!("exit with code: {:x?}", exit_code);
+        quit();
+        return;
+      }
     }
 
     self.write_mem(addr, self.dlen / 8, strobe, data);