From 865e0cc76689436c3ca0e487e208782c434676fd Mon Sep 17 00:00:00 2001
From: kunxian xia <xiakunxian130@gmail.com>
Date: Mon, 9 Sep 2024 22:18:57 +0800
Subject: [PATCH] add insn in step record when tracing

---
 ceno_emul/src/rv32im.rs   |  6 ++++++
 ceno_emul/src/tracer.rs   | 10 ++++++++++
 ceno_emul/src/vm_state.rs |  8 +++++++-
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/ceno_emul/src/rv32im.rs b/ceno_emul/src/rv32im.rs
index 710d7e7d4..64e758f55 100644
--- a/ceno_emul/src/rv32im.rs
+++ b/ceno_emul/src/rv32im.rs
@@ -183,6 +183,12 @@ pub struct Instruction {
     pub func7: u32,
 }
 
+impl Default for Instruction {
+    fn default() -> Self {
+        insn(InsnKind::INVALID, InsnCategory::Invalid, 0x00, 0x0, 0x00)
+    }
+}
+
 impl DecodedInstruction {
     pub fn new(insn: u32) -> Self {
         Self {
diff --git a/ceno_emul/src/tracer.rs b/ceno_emul/src/tracer.rs
index e2352ef0f..e8ae6efc5 100644
--- a/ceno_emul/src/tracer.rs
+++ b/ceno_emul/src/tracer.rs
@@ -5,6 +5,7 @@ use crate::{
     rv32im::DecodedInstruction,
     CENO_PLATFORM,
 };
+use crate::rv32im::Instruction;
 
 /// An instruction and its context in an execution trace. That is concrete values of registers and memory.
 ///
@@ -22,6 +23,7 @@ pub struct StepRecord {
     pub cycle: Cycle,
     pub pc: Change<ByteAddr>,
     pub insn_code: Word,
+    pub insn: Instruction,
 
     pub rs1: Option<ReadOp>,
     pub rs2: Option<ReadOp>,
@@ -62,6 +64,10 @@ impl StepRecord {
         DecodedInstruction::new(self.insn_code)
     }
 
+    pub fn insn(&self) -> Instruction {
+        self.insn
+    }
+
     pub fn rs1(&self) -> Option<ReadOp> {
         self.rs1.clone()
     }
@@ -134,6 +140,10 @@ impl Tracer {
         self.record.insn_code = value;
     }
 
+    pub fn store_insn(&mut self, insn: Instruction) {
+        self.record.insn = insn;
+    }
+
     pub fn load_register(&mut self, idx: RegIdx, value: Word) {
         let addr = CENO_PLATFORM.register_vma(idx).into();
 
diff --git a/ceno_emul/src/vm_state.rs b/ceno_emul/src/vm_state.rs
index b478fc5ff..546203fe8 100644
--- a/ceno_emul/src/vm_state.rs
+++ b/ceno_emul/src/vm_state.rs
@@ -82,6 +82,10 @@ impl VMState {
             Ok(step)
         }
     }
+
+    pub fn init_register_unsafe(&mut self, idx: RegIdx, value: Word) {
+        self.registers[idx] = value;
+    }
 }
 
 impl EmuContext for VMState {
@@ -109,7 +113,9 @@ impl EmuContext for VMState {
         Err(anyhow!("Trap {:?}", cause)) // Crash.
     }
 
-    fn on_insn_decoded(&mut self, _kind: &Instruction, _decoded: &DecodedInstruction) {}
+    fn on_insn_decoded(&mut self, insn: &Instruction, _decoded: &DecodedInstruction) {
+        self.tracer.store_insn(*insn);
+    }
 
     fn on_normal_end(&mut self, _kind: &Instruction, _decoded: &DecodedInstruction) {
         self.tracer.store_pc(ByteAddr(self.pc));