fix: Fixes for cycle marker and opcode stats (#11)

## What ❔

<!-- What are the changes this PR brings about? -->
<!-- Example: This PR adds a PR template to the repo. -->
<!-- (For bigger PRs adding more context is appreciated) -->

## Why ❔

<!-- Why are these changes done? What goal do they contribute to? What
are the principles behind them? -->
<!-- The `Why` has to be clear to non-Matter Labs entities running their
own ZK Chain -->
<!-- Example: PR templates ensure PR reviewers, observers, and future
iterators are in context about the evolution of repos. -->

## Is this a breaking change?
- [ ] Yes
- [ ] No

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [ ] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [ ] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [ ] Code has been formatted.

Author: antoniolocascio

risc_v_simulator/src/cycle/state.rs

@@ -28,7 +28,7 @@ pub const MARKER_CSR: u32 = 0x7ff;
 #[cfg(feature = "opcode_stats")]
 thread_local! {
 
-    static OPCODES_COUNTER: std::cell::RefCell<std::collections::HashMap<&'static str, usize>> = std::cell::RefCell::new(std::collections::HashMap::new());
+    pub(crate) static OPCODES_COUNTER: std::cell::RefCell<std::collections::HashMap<&'static str, usize>> = std::cell::RefCell::new(std::collections::HashMap::new());
 }
 
 #[inline(always)]
@@ -89,36 +89,42 @@ impl Mark {
 #[cfg(feature = "cycle_marker")]
 #[derive(Debug, Default)]
 pub struct CycleMarker {
+    cycle_counter: u64,
     pub markers: Vec<Mark>,
     pub delegation_counter: HashMap<u32, u64>,
 }
 
 #[cfg(feature = "cycle_marker")]
 impl CycleMarker {
-    fn new() -> Self {
+    pub(crate) fn new() -> Self {
         Self::default()
     }
 
     #[inline(always)]
-    fn add_marker(&mut self, count: u64) {
+    pub(crate) fn add_marker(&mut self) {
         self.markers.push(Mark {
-            cycles: count,
+            cycles: self.cycle_counter,
             delegations: self.delegation_counter.clone(),
         })
     }
 
     #[inline(always)]
-    fn add_delegation(&mut self, id: u32) {
+    pub(crate) fn add_delegation(&mut self, id: u32) {
         self.delegation_counter
             .entry(id)
             .and_modify(|n| *n += 1)
             .or_insert(1);
     }
+
+    #[inline(always)]
+    pub(crate) fn incr_cycle_counter(&mut self) {
+        self.cycle_counter += 1
+    }
 }
 
 #[cfg(feature = "cycle_marker")]
 thread_local! {
-  static CYCLE_MARKER: std::cell::RefCell<CycleMarker> = std::cell::RefCell::new(CycleMarker::new());
+  pub(crate) static CYCLE_MARKER: std::cell::RefCell<CycleMarker> = std::cell::RefCell::new(CycleMarker::new());
 }
 
 #[cfg(feature = "cycle_marker")]
@@ -264,7 +270,6 @@ pub struct RiscV32State<Config: MachineConfig = IMStandardIsaConfig> {
     pub pc: u32,
     pub extra_flags: ExtraFlags, // everything that doesn't need full register
 
-    pub cycle_counter: usize,
     pub timer: u64,
     pub timer_match: u64,
 
@@ -311,7 +316,6 @@ impl<Config: MachineConfig> RiscV32State<Config> {
             registers,
             pc,
             extra_flags,
-            cycle_counter,
             timer,
             timer_match,
             machine_mode_trap_data,
@@ -360,7 +364,6 @@ impl<Config: MachineConfig> RiscV32State<Config> {
             pc: initial_pc,
             extra_flags,
 
-            cycle_counter,
             timer,
             timer_match,
 
@@ -446,7 +449,7 @@ impl<Config: MachineConfig> RiscV32State<Config> {
     #[inline(always)]
     fn add_marker(&self) {
         #[cfg(feature = "cycle_marker")]
-        CYCLE_MARKER.with_borrow_mut(|cm| cm.add_marker(self.cycle_counter as u64))
+        CYCLE_MARKER.with_borrow_mut(|cm| cm.add_marker())
     }
 
     #[inline(always)]
@@ -455,6 +458,12 @@ impl<Config: MachineConfig> RiscV32State<Config> {
         CYCLE_MARKER.with_borrow_mut(|cm| cm.add_delegation(id))
     }
 
+    #[inline(always)]
+    fn count_new_cycle_for_markers(&self) {
+        #[cfg(feature = "cycle_marker")]
+        CYCLE_MARKER.with_borrow_mut(|cm| cm.incr_cycle_counter())
+    }
+
     pub fn cycle_ext<
         'a,
         M: MemorySource,
@@ -1394,10 +1403,7 @@ impl<Config: MachineConfig> RiscV32State<Config> {
 
         // Handle traps and interrupts.
         if trap.is_a_trap() {
-            println!(
-                "trap: {:?}, pc: {:08x}, proc_cycle: {:?}, instr: {:08x}",
-                trap, pc, self.cycle_counter, instr
-            );
+            println!("trap: {:?}, pc: {:08x}, instr: {:08x}", trap, pc, instr);
 
             if Config::HANDLE_EXCEPTIONS == false {
                 panic!("Simulator encountered an exception");
@@ -1440,7 +1446,7 @@ impl<Config: MachineConfig> RiscV32State<Config> {
 
         // for debugging
         self.sapt = mmu.read_sapt(current_privilege_mode, &mut trap);
-        self.cycle_counter += 1;
+        self.count_new_cycle_for_markers();
         tracer.at_cycle_end(&*self);
 
         //let trap = trap.as_register_value();

risc_v_simulator/src/cycle/state_new.rs

@@ -15,6 +15,10 @@ use crate::cycle::state::report_opcode;
 use crate::cycle::state::MARKER_CSR;
 use crate::cycle::state::NON_DETERMINISM_CSR;
 use crate::cycle::state::NUM_REGISTERS;
+#[cfg(feature = "opcode_stats")]
+use crate::cycle::state::OPCODES_COUNTER;
+#[cfg(feature = "cycle_marker")]
+use crate::cycle::state::{CycleMarker, Mark, CYCLE_MARKER};
 use crate::cycle::status_registers::TrapReason;
 use crate::cycle::IMStandardIsaConfig;
 use crate::cycle::MachineConfig;
@@ -148,7 +152,7 @@ impl<Config: MachineConfig> RiscV32StateForUnrolledProver<Config> {
     #[inline(always)]
     fn add_marker(&self) {
         #[cfg(feature = "cycle_marker")]
-        CYCLE_MARKER.with_borrow_mut(|cm| cm.add_marker(self.cycle_counter as u64))
+        CYCLE_MARKER.with_borrow_mut(|cm| cm.add_marker())
     }
 
     #[inline(always)]
@@ -157,6 +161,12 @@ impl<Config: MachineConfig> RiscV32StateForUnrolledProver<Config> {
         CYCLE_MARKER.with_borrow_mut(|cm| cm.add_delegation(id))
     }
 
+    #[inline(always)]
+    fn count_new_cycle_for_markers(&self) {
+        #[cfg(feature = "cycle_marker")]
+        CYCLE_MARKER.with_borrow_mut(|cm| cm.incr_cycle_counter())
+    }
+
     #[inline(always)]
     fn decoder_step<M: MemorySource, TR: Tracer<Config>>(
         &mut self,
@@ -854,6 +864,8 @@ impl<Config: MachineConfig> RiscV32StateForUnrolledProver<Config> {
                 }
             }
 
+            self.count_new_cycle_for_markers();
+
             tracer.at_cycle_end_ext(&*self);
 
             if self.pc == pc {

risc_v_simulator/src/runner/mod.rs

@@ -87,8 +87,10 @@ pub fn run_simple_for_num_cycles<S: NonDeterminismCSRSource<VectorMemoryImpl>, C
     memory.load_image(entry_point, binary.iter().copied());
 
     let mut previous_pc = entry_point;
+    let mut cycle_counter = 0u64;
 
     for _cycle in 0..cycles as usize {
+        cycle_counter += 1;
         RiscV32State::<C>::cycle(
             &mut state,
             &mut memory,
@@ -98,7 +100,7 @@ pub fn run_simple_for_num_cycles<S: NonDeterminismCSRSource<VectorMemoryImpl>, C
         );
 
         if state.pc == previous_pc {
-            println!("Took {} cycles to finish", state.cycle_counter);
+            println!("Took {} cycles to finish", cycle_counter);
             break;
         }
         previous_pc = state.pc;