test: Output memory size of each segment

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

@@ -125,7 +125,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
     }
 
     #[inline(always)]
-    pub fn check_and_segment(&mut self, instret: u64, segment_check_insns: u64) {
+    pub fn check_and_segment(&mut self, instret: u64, segment_check_insns: u64) -> bool {
         let threshold = self
             .segmentation_ctx
             .instret_last_segment_check
@@ -135,7 +135,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
             "overflow in segment check threshold calculation"
         );
         if instret < threshold {
-            return;
+            return false;
         }
 
         self.memory_ctx
@@ -149,6 +149,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
         if did_segment {
             self.reset_segment();
         }
+        did_segment
     }
 
     #[allow(dead_code)]
@@ -195,17 +196,36 @@ impl<const PAGE_BITS: usize> ExecutionCtxTrait for MeteredCtx<PAGE_BITS> {
     }
 
     #[inline(always)]
-    fn should_suspend<F>(
+    fn should_suspend<F: Clone>(
         instret: u64,
         _pc: u32,
         segment_check_insns: u64,
         exec_state: &mut VmExecState<F, GuestMemory, Self>,
     ) -> bool {
         // E2 always runs until termination. Here we use the function as a hook called every
         // instruction.
-        exec_state
+        let did_segment = exec_state
             .ctx
             .check_and_segment(instret, segment_check_insns);
+        if did_segment {
+            let num_fs = exec_state.vm_state.streams.hint_stream.len()
+                + exec_state
+                    .vm_state
+                    .streams
+                    .input_stream
+                    .iter()
+                    .map(|x| x.len())
+                    .sum::<usize>()
+                + exec_state.vm_state.streams.hint_space.len();
+            let num_pages = exec_state.ctx.memory_ctx.global_page_access_count;
+            let segment_idx = exec_state.ctx.segmentation_ctx.segments.len() - 1;
+            tracing::info!(
+                "Segment idx: {}, Memory size: {}bytes, stream size {}bytes",
+                segment_idx,
+                num_pages * (1 << PAGE_BITS),
+                num_fs * 4
+            );
+        }
         false
     }
 

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

@@ -14,6 +14,11 @@ impl BitSet {
         }
     }
 
+    pub fn num_of_ones(&self) -> usize {
+        let num_ones: u32 = self.words.iter().map(|x| x.count_ones()).sum();
+        num_ones as usize
+    }
+
     #[inline(always)]
     pub fn insert(&mut self, index: usize) -> bool {
         let word_index = index >> 6;
@@ -99,6 +104,7 @@ impl BitSet {
 #[derive(Clone, Debug)]
 pub struct MemoryCtx<const PAGE_BITS: usize> {
     pub page_indices: BitSet,
+    pub global_page_indices: BitSet,
     memory_dimensions: MemoryDimensions,
     min_block_size_bits: Vec<u8>,
     pub boundary_idx: usize,
@@ -108,6 +114,7 @@ pub struct MemoryCtx<const PAGE_BITS: usize> {
     chunk: u32,
     chunk_bits: u32,
     page_access_count: usize,
+    pub global_page_access_count: usize,
     // Note: 32 is the maximum access adapter size.
     addr_space_access_count: Vec<usize>,
 }
@@ -123,6 +130,7 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
         Self {
             // Address height already considers `chunk_bits`.
             page_indices: BitSet::new(1 << (merkle_height.saturating_sub(PAGE_BITS))),
+            global_page_indices: BitSet::new(1 << (merkle_height.saturating_sub(PAGE_BITS))),
             min_block_size_bits: config.memory_config.min_block_size_bits(),
             boundary_idx: config.memory_boundary_air_id(),
             merkle_tree_index: config.memory_merkle_air_id(),
@@ -132,6 +140,7 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
             memory_dimensions,
             continuations_enabled: config.continuation_enabled,
             page_access_count: 0,
+            global_page_access_count: 0,
             addr_space_access_count: vec![0; (1 << memory_dimensions.addr_space_height) + 1],
         }
     }
@@ -178,6 +187,9 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
         let end_page_id = ((end_block_id - 1) >> PAGE_BITS) + 1;
 
         for page_id in start_page_id..end_page_id {
+            if self.global_page_indices.insert(page_id as usize) {
+                self.global_page_access_count += 1;
+            }
             if self.page_indices.insert(page_id as usize) {
                 self.page_access_count += 1;
                 // SAFETY: address_space passed is usually a hardcoded constant or derived from an

crates/vm/src/arch/execution_mode/mod.rs

@@ -13,7 +13,7 @@ pub use pure::ExecutionCtx;
 pub trait ExecutionCtxTrait: Sized {
     fn on_memory_operation(&mut self, address_space: u32, ptr: u32, size: u32);
 
-    fn should_suspend<F>(
+    fn should_suspend<F: Clone>(
         instret: u64,
         pc: u32,
         _arg: u64,

crates/vm/src/arch/state.rs

@@ -17,7 +17,7 @@ use crate::{
 };
 
 /// Represents the core state of a VM.
-#[derive(derive_new::new, CopyGetters, MutGetters)]
+#[derive(derive_new::new, CopyGetters, MutGetters, Clone)]
 pub struct VmState<F, MEM = GuestMemory> {
     #[getset(get_copy = "pub", get_mut = "pub")]
     instret: u64,