fix: avoid stack push/pop

Author: shuklaayush

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

@@ -190,6 +190,15 @@ 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 {
+            // Append page_id to page_indices_since_checkpoint
+            let len = self.page_indices_since_checkpoint_len;
+            debug_assert!(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) {
                 // SAFETY: address_space passed is usually a hardcoded constant or derived from an
                 // Instruction where it is bounds checked before passing
@@ -199,15 +208,6 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
                         .get_unchecked_mut(address_space as usize) += 1;
                 }
             }
-
-            // Append page_id to page_indices_since_checkpoint
-            let len = self.page_indices_since_checkpoint_len;
-            debug_assert!(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;
         }
     }
 

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

@@ -288,8 +288,22 @@ mod aot {
         // }
         // self.page_indices_since_checkpoint.push(page_id);
 
-        // Start with `ptr_reg = index`
-        asm_str += &format!("    push {ptr_reg}\n");
+        // Append page_id to page_indices_since_checkpoint
+        asm_str += &format!(
+            "    mov {reg1}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}]\n"
+        );
+        asm_str += &format!(
+            "    mov {reg2}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_ptr_offset}]\n"
+        );
+        let ptr_reg_32 = convert_x86_reg(ptr_reg, Width::W32).ok_or_else(|| {
+            AotError::Other(format!("unsupported ptr_reg for 32-bit store: {ptr_reg}"))
+        })?;
+        asm_str += &format!("    mov dword ptr [{reg2} + {reg1} * 4], {ptr_reg_32}\n");
+        asm_str += &format!("    add {reg1}, 1\n");
+        asm_str += &format!(
+            "    mov [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}], {reg1}\n"
+        );
+
         // `reg1 = word_index`
         asm_str += &format!("    mov {reg1}, {ptr_reg}\n");
         asm_str += &format!("    shr {reg1}, 6\n");
@@ -322,22 +336,6 @@ 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");
-        // Append page_id to page_indices_since_checkpoint
-        asm_str += &format!("    pop {ptr_reg}\n");
-        asm_str += &format!(
-            "    mov {reg1}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}]\n"
-        );
-        asm_str += &format!(
-            "    mov {reg2}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_ptr_offset}]\n"
-        );
-        let ptr_reg_32 = convert_x86_reg(ptr_reg, Width::W32).ok_or_else(|| {
-            AotError::Other(format!("unsupported ptr_reg for 32-bit store: {ptr_reg}"))
-        })?;
-        asm_str += &format!("    mov dword ptr [{reg2} + {reg1} * 4], {ptr_reg_32}\n");
-        asm_str += &format!("    add {reg1}, 1\n");
-        asm_str += &format!(
-            "    mov [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}], {reg1}\n"
-        );
 
         Ok(asm_str)
     }