perf: Use u32 for address_space and address internally (#1101)

* perf: Use usize for address_space internally everywhere

* u32 instead of usize

* Comments

crates/sdk/src/verifier/leaf/mod.rs

@@ -112,9 +112,7 @@ impl LeafVmVerifierConfig {
         builder: &mut Builder<C>,
     ) -> ([Felt<F>; DIGEST_SIZE], [Felt<F>; DIGEST_SIZE]) {
         let memory_dimensions = self.app_system_config.memory_config.memory_dimensions();
-        let pv_as = F::from_canonical_usize(
-            PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset,
-        );
+        let pv_as = PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset;
         let pv_start_idx = memory_dimensions.label_to_index((pv_as, 0));
         let pv_height = log2_strict_usize(self.app_system_config.num_public_values / DIGEST_SIZE);
         let proof_len = memory_dimensions.overall_height() - pv_height;

crates/toolchain/instructions/src/exe.rs

@@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize};
 use crate::program::Program;
 
 /// Memory image is a map from (address space, address) to word.
-pub type MemoryImage<F> = BTreeMap<(F, F), F>;
+pub type MemoryImage<F> = BTreeMap<(u32, u32), F>;
 /// Stores the starting address, end address, and name of a set of function.
 pub type FnBounds = BTreeMap<u32, FnBound>;
 

crates/toolchain/transpiler/src/util.rs

@@ -1,8 +1,11 @@
 use std::collections::BTreeMap;
 
 use openvm_instructions::{
-    exe::MemoryImage, instruction::Instruction, riscv::RV32_REGISTER_NUM_LIMBS,
-    utils::isize_to_field, SystemOpcode, VmOpcode,
+    exe::MemoryImage,
+    instruction::Instruction,
+    riscv::{RV32_MEMORY_AS, RV32_REGISTER_NUM_LIMBS},
+    utils::isize_to_field,
+    SystemOpcode, VmOpcode,
 };
 use openvm_stark_backend::p3_field::PrimeField32;
 use rrs_lib::instruction_formats::{BType, IType, ITypeShamt, JType, RType, SType, UType};
@@ -167,7 +170,7 @@ pub fn elf_memory_image_to_openvm_memory_image<F: PrimeField32>(
     for (addr, word) in memory_image {
         for (i, byte) in word.to_le_bytes().into_iter().enumerate() {
             result.insert(
-                (F::TWO, F::from_canonical_u32(addr + i as u32)),
+                (RV32_MEMORY_AS, addr + i as u32),
                 F::from_canonical_u8(byte),
             );
         }

crates/vm/src/arch/config.rs

@@ -47,7 +47,7 @@ pub struct MemoryConfig {
     /// The maximum height of the address space. This means the trie has `as_height` layers for searching the address space. The allowed address spaces are those in the range `[as_offset, as_offset + 2^as_height)` where `as_offset` is currently fixed to `1` to not allow address space `0` in memory.
     pub as_height: usize,
     /// The offset of the address space.
-    pub as_offset: usize,
+    pub as_offset: u32,
     pub pointer_max_bits: usize,
     pub clk_max_bits: usize,
     /// Limb size used by the range checker

crates/vm/src/system/memory/manager/dimensions.rs

@@ -1,5 +1,5 @@
 use derive_new::new;
-use openvm_stark_backend::{p3_field::PrimeField32, p3_util::log2_strict_usize};
+use openvm_stark_backend::p3_util::log2_strict_usize;
 
 use crate::{arch::MemoryConfig, system::memory::CHUNK};
 
@@ -12,18 +12,17 @@ pub struct MemoryDimensions {
     /// Pointer height
     pub address_height: usize,
     /// Address space offset
-    pub as_offset: usize,
+    pub as_offset: u32,
 }
 
 impl MemoryDimensions {
     pub fn overall_height(&self) -> usize {
         self.as_height + self.address_height
     }
     /// Convert an address label (address space, block id) to its index in the memory merkle tree.
-    pub fn label_to_index<F: PrimeField32>(&self, label: (F, usize)) -> usize {
+    pub fn label_to_index(&self, label: (u32, u32)) -> u64 {
         let (addr_space, block_id) = label;
-        ((addr_space.as_canonical_u32() as usize - self.as_offset) << self.address_height)
-            + block_id
+        (((addr_space - self.as_offset) as u64) << self.address_height) + block_id as u64
     }
 }
 

crates/vm/src/system/memory/manager/interface.rs

@@ -21,7 +21,7 @@ pub enum MemoryInterface<F> {
 }
 
 impl<F: PrimeField32> MemoryInterface<F> {
-    pub fn touch_address(&mut self, addr_space: F, pointer: F) {
+    pub fn touch_address(&mut self, addr_space: u32, pointer: u32) {
         match self {
             MemoryInterface::Volatile { boundary_chip } => {
                 boundary_chip.touch_address(addr_space, pointer);