fix(new-execution): add register contribtions on reset_segment (#1843)

for metering, we assume all registers are going to be read during
execution and don't add their contributions in the hot loop.
this pr fixes the `reset_segment` function to also add the register
contributions

Author: shuklaayush

crates/vm/src/arch/execution_mode/metered/ctx.rs

@@ -15,6 +15,7 @@ pub struct MeteredCtx<const PAGE_BITS: usize = 6> {
 
     pub memory_ctx: MemoryCtx<PAGE_BITS>,
     pub segmentation_ctx: SegmentationCtx,
+    pub continuations_enabled: bool,
 }
 
 impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
@@ -29,7 +30,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
         widths: Vec<usize>,
         interactions: Vec<usize>,
     ) -> Self {
-        let (mut trace_heights, is_trace_height_constant): (Vec<u32>, Vec<bool>) =
+        let (trace_heights, is_trace_height_constant): (Vec<u32>, Vec<bool>) =
             constant_trace_heights
                 .iter()
                 .map(|&constant_height| {
@@ -41,7 +42,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
                 })
                 .unzip();
 
-        let mut memory_ctx = MemoryCtx::new(
+        let memory_ctx = MemoryCtx::new(
             has_public_values_chip,
             continuations_enabled,
             as_byte_alignment_bits,
@@ -57,22 +58,29 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
 
         let segmentation_ctx = SegmentationCtx::new(air_names, widths, interactions);
 
+        let mut ctx = Self {
+            trace_heights,
+            is_trace_height_constant,
+            memory_ctx,
+            segmentation_ctx,
+            continuations_enabled,
+        };
+
         // Add merkle height contributions for all registers
-        if continuations_enabled {
-            memory_ctx.update_boundary_merkle_heights(
-                &mut trace_heights,
+        ctx.add_register_merkle_heights();
+
+        ctx
+    }
+
+    fn add_register_merkle_heights(&mut self) {
+        if self.continuations_enabled {
+            self.memory_ctx.update_boundary_merkle_heights(
+                &mut self.trace_heights,
                 RV32_REGISTER_AS,
                 0,
                 (RV32_NUM_REGISTERS * RV32_REGISTER_NUM_LIMBS) as u32,
             );
         }
-
-        Self {
-            trace_heights,
-            is_trace_height_constant,
-            memory_ctx,
-            segmentation_ctx,
-        }
     }
 
     pub fn with_max_trace_height(mut self, max_trace_height: u32) -> Self {
@@ -111,6 +119,9 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
                 self.trace_heights[i] = 0;
             }
         }
+
+        // Add merkle height contributions for all registers
+        self.add_register_merkle_heights();
     }
 
     pub fn check_and_segment(&mut self, instret: u64) {
@@ -124,6 +135,21 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
             self.reset_segment();
         }
     }
+
+    #[allow(dead_code)]
+    pub fn print_heights(&self) {
+        println!("{:>10} {:<30}", "Height", "Air Name");
+        println!("{}", "-".repeat(42));
+        for (i, height) in self.trace_heights.iter().enumerate() {
+            let air_name = self
+                .segmentation_ctx
+                .air_names
+                .get(i)
+                .map(|s| s.as_str())
+                .unwrap_or("Unknown");
+            println!("{:>10} {:<30}", height, air_name);
+        }
+    }
 }
 
 impl<const PAGE_BITS: usize> E1E2ExecutionCtx for MeteredCtx<PAGE_BITS> {

crates/vm/src/arch/execution_mode/metered/segment_ctx.rs

@@ -37,7 +37,7 @@ impl Default for SegmentationLimits {
 pub struct SegmentationCtx {
     pub segments: Vec<Segment>,
     instret_last_segment_check: u64,
-    air_names: Vec<String>,
+    pub(crate) air_names: Vec<String>,
     widths: Vec<usize>,
     interactions: Vec<usize>,
     segment_check_insns: u64,