chore: clean up

Author: shuklaayush

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

@@ -107,11 +107,10 @@ pub struct MemoryCtx<const PAGE_BITS: usize> {
     pub boundary_idx: usize,
     pub merkle_tree_index: Option<usize>,
     pub adapter_offset: usize,
-    pub continuations_enabled: bool,
+    continuations_enabled: bool,
     chunk: u32,
     chunk_bits: u32,
     pub page_indices: BitSet,
-    pub page_access_count: usize,
     pub addr_space_access_count: RVec<usize>,
     pub page_indices_since_checkpoint: Box<[u32]>,
     pub page_indices_since_checkpoint_len: usize,
@@ -140,7 +139,6 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
             memory_dimensions,
             continuations_enabled: config.continuation_enabled,
             page_indices: BitSet::new(bitset_size),
-            page_access_count: 0,
             addr_space_access_count: vec![0; addr_space_size].into(),
             page_indices_since_checkpoint: vec![0; page_indices_since_checkpoint_cap]
                 .into_boxed_slice(),
@@ -257,7 +255,10 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
     /// Initialize state for a new segment
     #[inline(always)]
     pub(crate) fn initialize_segment(&mut self, trace_heights: &mut [u32]) {
-        // Reset trace heights for memory chips
+        // Clear page indices for the new segment
+        self.page_indices.clear();
+
+        // Reset trace heights for memory chips as 0
         // SAFETY: boundary_idx is a compile time constant within bounds
         unsafe {
             *trace_heights.get_unchecked_mut(self.boundary_idx) = 0;
@@ -277,15 +278,11 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
         // Apply height updates for all pages accessed since last checkpoint, and
         // initialize page_indices for the new segment.
         let mut addr_space_access_count = vec![0; self.addr_space_access_count.len()];
-        let mut page_access_count = 0;
-        self.page_indices.clear();
-
         let pages_len = self.page_indices_since_checkpoint_len;
         for i in 0..pages_len {
             // SAFETY: i is within 0..pages_len and pages_len is the slice length.
             let page_id = unsafe { *self.page_indices_since_checkpoint.get_unchecked(i) } as usize;
             if self.page_indices.insert(page_id) {
-                page_access_count += 1;
                 let (addr_space, _) = self
                     .memory_dimensions
                     .index_to_label((page_id as u64) << PAGE_BITS);
@@ -298,9 +295,7 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
                 }
             }
         }
-
-        self.apply_height_updates(trace_heights, page_access_count, &addr_space_access_count);
-
+        self.apply_height_updates(trace_heights, &addr_space_access_count);
         self.page_indices_since_checkpoint_len = 0;
     }
 
@@ -312,12 +307,9 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
 
     /// Apply height updates given page counts
     #[inline(always)]
-    fn apply_height_updates(
-        &self,
-        trace_heights: &mut [u32],
-        page_access_count: usize,
-        addr_space_access_count: &[usize],
-    ) {
+    fn apply_height_updates(&self, trace_heights: &mut [u32], addr_space_access_count: &[usize]) {
+        let page_access_count = addr_space_access_count.iter().sum();
+
         // On page fault, assume we add all leaves in a page
         let leaves = (page_access_count << PAGE_BITS) as u32;
         // SAFETY: boundary_idx is a compile time constant within bounds
@@ -346,7 +338,9 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
             }
         }
 
-        for (address_space, &x) in addr_space_access_count.iter().enumerate() {
+        for address_space in 0..addr_space_access_count.len() {
+            // SAFETY: address_space is from 0 to len(), guaranteed to be in bounds
+            let x = unsafe { *addr_space_access_count.get_unchecked(address_space) };
             if x > 0 {
                 // Initial **and** final handling of touched pages requires send (resp. receive) in
                 // chunk-sized units for the merkle chip
@@ -365,13 +359,7 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
     /// Resolve all lazy updates of each memory access for memory adapters/poseidon2/merkle chip.
     #[inline(always)]
     pub(crate) fn lazy_update_boundary_heights(&mut self, trace_heights: &mut [u32]) {
-        let page_access_count = self.addr_space_access_count.iter().sum();
-        self.apply_height_updates(
-            trace_heights,
-            page_access_count,
-            &self.addr_space_access_count,
-        );
-        self.page_access_count = 0;
+        self.apply_height_updates(trace_heights, &self.addr_space_access_count);
         // SAFETY: Resetting array elements to 0 is always safe
         unsafe {
             std::ptr::write_bytes(

extensions/rv32im/circuit/src/common/mod.rs

@@ -212,8 +212,15 @@ mod aot {
         //     let end_page_id = ((end_block_id - 1) >> PAGE_BITS) + 1;
 
         //     for page_id in start_page_id..end_page_id {
+        //          // Append page_id to page_indices_since_checkpoint
+        //          let len = self.page_indices_since_checkpoint_len;
+        //          // SAFETY: len is within bounds, and we extend length by 1 after writing.
+        //          unsafe {
+        //              *self.page_indices_since_checkpoint.as_mut_ptr().add(len) = page_id;
+        //          }
+        //          self.page_indices_since_checkpoint_len = len + 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             // Instruction where it is bounds checked before passing
         //             unsafe {
@@ -276,17 +283,6 @@ mod aot {
                 page_indices_since_checkpoint_len
             );
         let inserted_label = format!(".asm_execute_pc_{pc}_inserted");
-        // The next section is the implementation of `BitSet::insert` in ASM.
-        // pub fn insert(&mut self, index: usize) -> bool {
-        //     let word_index = index >> 6;
-        //     let bit_index = index & 63;
-        //     let mask = 1u64 << bit_index;
-        //     let word = unsafe { self.words.get_unchecked_mut(word_index) };
-        //     let was_set = (*word & mask) != 0;
-        //     *word |= mask;
-        //     !was_set
-        // }
-        // self.page_indices_since_checkpoint.push(page_id);
 
         // Append page_id to page_indices_since_checkpoint
         asm_str += &format!(
@@ -304,6 +300,18 @@ mod aot {
             "    mov [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}], {reg1}\n"
         );
 
+        // The next section is the implementation of `BitSet::insert` in ASM.
+        // pub fn insert(&mut self, index: usize) -> bool {
+        //     let word_index = index >> 6;
+        //     let bit_index = index & 63;
+        //     let mask = 1u64 << bit_index;
+        //     let word = unsafe { self.words.get_unchecked_mut(word_index) };
+        //     let was_set = (*word & mask) != 0;
+        //     *word |= mask;
+        //     !was_set
+        // }
+
+        // Start with `ptr_reg = index`
         // `reg1 = word_index`
         asm_str += &format!("    mov {reg1}, {ptr_reg}\n");
         asm_str += &format!("    shr {reg1}, 6\n");
@@ -336,6 +344,7 @@ mod aot {
         // self.addr_space_access_count[address_space] += 1;
         asm_str += &format!("    add dword ptr [{reg1} + {address_space} * 4], 1\n");
         asm_str += &format!("{inserted_label}:\n");
+        // Inserted, do nothing
 
         Ok(asm_str)
     }