diff --git a/crates/sdk/src/verifier/leaf/mod.rs b/crates/sdk/src/verifier/leaf/mod.rs index 5cc1cfa113..71659bbc3d 100644 --- a/crates/sdk/src/verifier/leaf/mod.rs +++ b/crates/sdk/src/verifier/leaf/mod.rs @@ -112,9 +112,7 @@ impl LeafVmVerifierConfig { builder: &mut Builder, ) -> ([Felt; DIGEST_SIZE], [Felt; 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; diff --git a/crates/toolchain/instructions/src/exe.rs b/crates/toolchain/instructions/src/exe.rs index 9e5299a22f..fb84ec7da5 100644 --- a/crates/toolchain/instructions/src/exe.rs +++ b/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 = BTreeMap<(F, F), F>; +pub type MemoryImage = BTreeMap<(u32, u32), F>; /// Stores the starting address, end address, and name of a set of function. pub type FnBounds = BTreeMap; diff --git a/crates/toolchain/transpiler/src/util.rs b/crates/toolchain/transpiler/src/util.rs index 8e5cfcefa9..b1f8a24e7c 100644 --- a/crates/toolchain/transpiler/src/util.rs +++ b/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( 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), ); } diff --git a/crates/vm/src/arch/config.rs b/crates/vm/src/arch/config.rs index 715a0f1542..41a926d92b 100644 --- a/crates/vm/src/arch/config.rs +++ b/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 diff --git a/crates/vm/src/system/memory/manager/dimensions.rs b/crates/vm/src/system/memory/manager/dimensions.rs index e91fe86988..ebe3612198 100644 --- a/crates/vm/src/system/memory/manager/dimensions.rs +++ b/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,7 +12,7 @@ pub struct MemoryDimensions { /// Pointer height pub address_height: usize, /// Address space offset - pub as_offset: usize, + pub as_offset: u32, } impl MemoryDimensions { @@ -20,10 +20,9 @@ impl MemoryDimensions { 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(&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 } } diff --git a/crates/vm/src/system/memory/manager/interface.rs b/crates/vm/src/system/memory/manager/interface.rs index a3a69d8b1a..a950afdb45 100644 --- a/crates/vm/src/system/memory/manager/interface.rs +++ b/crates/vm/src/system/memory/manager/interface.rs @@ -21,7 +21,7 @@ pub enum MemoryInterface { } impl MemoryInterface { - 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); diff --git a/crates/vm/src/system/memory/manager/memory.rs b/crates/vm/src/system/memory/manager/memory.rs index ba9f4c176c..c38016219f 100644 --- a/crates/vm/src/system/memory/manager/memory.rs +++ b/crates/vm/src/system/memory/manager/memory.rs @@ -48,11 +48,11 @@ impl MemoryReadRecord { pub const INITIAL_TIMESTAMP: u32 = 0; /// (address_space, pointer) -type Address = (usize, usize); +type Address = (u32, u32); #[derive(Clone, Copy, PartialEq, Eq, Debug)] struct BlockData { - pointer: usize, + pointer: u32, size: usize, timestamp: u32, } @@ -76,16 +76,15 @@ impl Memory { let mut block_data = FxHashMap::default(); let mut data = FxHashMap::default(); for (&(address_space, block_idx), values) in initial_memory { - let address_space_usize = address_space.as_canonical_u32() as usize; - let pointer = block_idx * N; + let pointer = block_idx * N as u32; let block = BlockData { pointer, size: N, timestamp: INITIAL_TIMESTAMP, }; for (i, value) in values.iter().enumerate() { - data.insert((address_space_usize, pointer + i), *value); - block_data.insert((address_space_usize, pointer + i), block); + data.insert((address_space, pointer + i as u32), *value); + block_data.insert((address_space, pointer + i as u32), block); } } Self { @@ -113,8 +112,8 @@ impl Memory { /// Writes an array of values to the memory at the specified address space and start index. pub fn write( &mut self, - address_space: usize, - pointer: usize, + address_space: u32, + pointer: u32, values: [F; N], ) -> (MemoryWriteRecord, Vec>) { assert!(N.is_power_of_two()); @@ -127,13 +126,13 @@ impl Memory { let prev_data = array::from_fn(|i| { self.data - .insert((address_space, pointer + i), values[i]) + .insert((address_space, pointer + i as u32), values[i]) .unwrap_or(F::ZERO) }); let record = MemoryWriteRecord { - address_space: F::from_canonical_usize(address_space), - pointer: F::from_canonical_usize(pointer), + address_space: F::from_canonical_u32(address_space), + pointer: F::from_canonical_u32(pointer), timestamp: self.timestamp, prev_timestamp, data: values, @@ -146,8 +145,8 @@ impl Memory { /// Reads an array of values from the memory at the specified address space and start index. pub fn read( &mut self, - address_space: usize, - pointer: usize, + address_space: u32, + pointer: u32, ) -> (MemoryReadRecord, Vec>) { assert!(N.is_power_of_two()); @@ -158,8 +157,8 @@ impl Memory { debug_assert!(prev_timestamp < self.timestamp); let record = MemoryReadRecord { - address_space: F::from_canonical_usize(address_space), - pointer: F::from_canonical_usize(pointer), + address_space: F::from_canonical_u32(address_space), + pointer: F::from_canonical_u32(pointer), timestamp: self.timestamp, prev_timestamp, data: self.range_array::(address_space, pointer), @@ -179,7 +178,7 @@ impl Memory { let to_access: FxHashSet<_> = self .block_data .keys() - .map(|&(address_space, pointer)| (address_space, (pointer / N) * N)) + .map(|&(address_space, pointer)| (address_space, (pointer / N as u32) * N as u32)) .collect(); for &(address_space, pointer) in to_access.iter() { @@ -193,11 +192,11 @@ impl Memory { for (address_space, pointer) in to_access { let block = self.block_data.get(&(address_space, pointer)).unwrap(); - debug_assert_eq!(block.pointer % N, 0); + debug_assert_eq!(block.pointer % N as u32, 0); debug_assert_eq!(block.size, N); equipartition.insert( - (F::from_canonical_usize(address_space), pointer / N), + (address_space, pointer / N as u32), TimestampedValues { timestamp: block.timestamp, values: self.range_array::(address_space, pointer), @@ -211,8 +210,8 @@ impl Memory { // Modifies the partition to ensure that there is a block starting at (address_space, query). fn split_to_make_boundary( &mut self, - address_space: usize, - query: usize, + address_space: u32, + query: u32, records: &mut Vec>, ) { let original_block = self.block_containing(address_space, query); @@ -230,44 +229,43 @@ impl Memory { // Split. records.push(AccessAdapterRecord { timestamp, - address_space: F::from_canonical_usize(address_space), - start_index: F::from_canonical_usize(cur_ptr), - data: data - [cur_ptr - original_block.pointer..cur_ptr - original_block.pointer + cur_size] + address_space: F::from_canonical_u32(address_space), + start_index: F::from_canonical_u32(cur_ptr), + data: data[(cur_ptr - original_block.pointer) as usize + ..(cur_ptr - original_block.pointer) as usize + cur_size] .to_vec(), kind: AccessAdapterRecordKind::Split, }); let half_size = cur_size / 2; + let half_size_u32 = half_size as u32; + let mid_ptr = cur_ptr + half_size_u32; - if query <= cur_ptr + half_size { + if query <= mid_ptr { // The right is finalized; add it to the partition. let block = BlockData { - pointer: cur_ptr + half_size, + pointer: mid_ptr, size: half_size, timestamp, }; - for i in 0..half_size { - self.block_data - .insert((address_space, cur_ptr + half_size + i), block); + for i in 0..half_size_u32 { + self.block_data.insert((address_space, mid_ptr + i), block); } } - if query >= cur_ptr + half_size { + if query >= cur_ptr + half_size_u32 { // The left is finalized; add it to the partition. let block = BlockData { pointer: cur_ptr, size: half_size, timestamp, }; - for i in 0..half_size { + for i in 0..half_size_u32 { self.block_data.insert((address_space, cur_ptr + i), block); } } - - if cur_ptr + half_size <= query { - cur_ptr += half_size; + if mid_ptr <= query { + cur_ptr = mid_ptr; } - if cur_ptr == query { break; } @@ -277,8 +275,8 @@ impl Memory { fn access_updating_timestamp( &mut self, - address_space: usize, - pointer: usize, + address_space: u32, + pointer: u32, size: usize, records: &mut Vec>, ) -> u32 { @@ -286,7 +284,7 @@ impl Memory { let mut prev_timestamp = None; - for i in 0..size { + for i in 0..size as u32 { let block = self .block_data .get_mut(&(address_space, pointer + i)) @@ -300,13 +298,13 @@ impl Memory { fn access( &mut self, - address_space: usize, - pointer: usize, + address_space: u32, + pointer: u32, size: usize, records: &mut Vec>, ) { self.split_to_make_boundary(address_space, pointer, records); - self.split_to_make_boundary(address_space, pointer + size, records); + self.split_to_make_boundary(address_space, pointer + size as u32, records); let block_data = self .block_data @@ -315,8 +313,8 @@ impl Memory { .unwrap_or_else(|| { for i in 0..size { self.block_data.insert( - (address_space, pointer + i), - self.initial_block_data(pointer + i), + (address_space, pointer + i as u32), + self.initial_block_data(pointer + i as u32), ); } self.initial_block_data(pointer) @@ -330,7 +328,12 @@ impl Memory { // Now recursively access left and right blocks to ensure they are in the partition. let half_size = size / 2; self.access(address_space, pointer, half_size, records); - self.access(address_space, pointer + half_size, half_size, records); + self.access( + address_space, + pointer + half_size as u32, + half_size, + records, + ); self.merge_block_with_next(address_space, pointer, records); } @@ -341,8 +344,8 @@ impl Memory { /// do not have the same size. fn merge_block_with_next( &mut self, - address_space: usize, - pointer: usize, + address_space: u32, + pointer: u32, records: &mut Vec>, ) { let left_block = self.block_data.get(&(address_space, pointer)).unwrap(); @@ -352,12 +355,12 @@ impl Memory { let right_timestamp = self .block_data - .get(&(address_space, pointer + size)) + .get(&(address_space, pointer + size as u32)) .map(|b| b.timestamp) .unwrap_or(INITIAL_TIMESTAMP); let timestamp = max(left_timestamp, right_timestamp); - for i in 0..2 * size { + for i in 0..2 * size as u32 { self.block_data.insert( (address_space, pointer + i), BlockData { @@ -369,8 +372,8 @@ impl Memory { } records.push(AccessAdapterRecord { timestamp, - address_space: F::from_canonical_usize(address_space), - start_index: F::from_canonical_usize(pointer), + address_space: F::from_canonical_u32(address_space), + start_index: F::from_canonical_u32(pointer), data: self.range_vec(address_space, pointer, 2 * size), kind: AccessAdapterRecordKind::Merge { left_timestamp, @@ -379,7 +382,7 @@ impl Memory { }); } - fn block_containing(&mut self, address_space: usize, pointer: usize) -> BlockData { + fn block_containing(&mut self, address_space: u32, pointer: u32) -> BlockData { if let Some(block_data) = self.block_data.get(&(address_space, pointer)) { *block_data } else { @@ -387,8 +390,9 @@ impl Memory { } } - fn initial_block_data(&self, pointer: usize) -> BlockData { - let aligned_pointer = (pointer / self.initial_block_size) * self.initial_block_size; + fn initial_block_data(&self, pointer: u32) -> BlockData { + let aligned_pointer = + (pointer / self.initial_block_size as u32) * self.initial_block_size as u32; BlockData { pointer: aligned_pointer, size: self.initial_block_size, @@ -396,17 +400,17 @@ impl Memory { } } - pub fn get(&self, address_space: usize, pointer: usize) -> F { + pub fn get(&self, address_space: u32, pointer: u32) -> F { *self.data.get(&(address_space, pointer)).unwrap_or(&F::ZERO) } - fn range_array(&self, address_space: usize, pointer: usize) -> [F; N] { - array::from_fn(|i| self.get(address_space, pointer + i)) + fn range_array(&self, address_space: u32, pointer: u32) -> [F; N] { + array::from_fn(|i| self.get(address_space, pointer + i as u32)) } - fn range_vec(&self, address_space: usize, pointer: usize, len: usize) -> Vec { + fn range_vec(&self, address_space: u32, pointer: u32, len: usize) -> Vec { (0..len) - .map(|i| self.get(address_space, pointer + i)) + .map(|i| self.get(address_space, pointer + i as u32)) .collect() } } @@ -811,7 +815,7 @@ mod tests { let (final_memory, records) = memory.finalize::<8>(); assert_eq!(final_memory.len(), 4); assert_eq!( - final_memory.get(&(bb!(1), 0)), + final_memory.get(&(1, 0)), Some(&TimestampedValues { values: bba![1, 2, 3, 4, 0, 0, 0, 0], timestamp: 1, @@ -819,7 +823,7 @@ mod tests { ); // start_index = 16 corresponds to label = 2 assert_eq!( - final_memory.get(&(bb!(1), 2)), + final_memory.get(&(1, 2)), Some(&TimestampedValues { values: bba![1, 1, 1, 1, 1, 1, 1, 1], timestamp: 2, @@ -827,7 +831,7 @@ mod tests { ); // start_index = 24 corresponds to label = 3 assert_eq!( - final_memory.get(&(bb!(1), 3)), + final_memory.get(&(1, 3)), Some(&TimestampedValues { values: bba![1, 1, 1, 1, 1, 1, 1, 1], timestamp: 2, @@ -835,7 +839,7 @@ mod tests { ); // start_index = 64 corresponds to label = 8 assert_eq!( - final_memory.get(&(bb!(2), 8)), + final_memory.get(&(2, 8)), Some(&TimestampedValues { values: bba![8, 7, 6, 5, 4, 3, 2, 1], timestamp: 3, @@ -852,8 +856,8 @@ mod tests { // Initialize initial memory with blocks at indices 0 and 2 let mut initial_memory = Equipartition::::new(); - initial_memory.insert((F::ONE, 0), bba![1, 2, 3, 4, 5, 6, 7, 8]); // Block 0, pointers 0–8 - initial_memory.insert((F::ONE, 2), bba![1, 2, 3, 4, 5, 6, 7, 8]); // Block 2, pointers 16–24 + initial_memory.insert((1, 0), bba![1, 2, 3, 4, 5, 6, 7, 8]); // Block 0, pointers 0–8 + initial_memory.insert((1, 2), bba![1, 2, 3, 4, 5, 6, 7, 8]); // Block 2, pointers 16–24 let mut memory = Memory::new(&initial_memory); diff --git a/crates/vm/src/system/memory/manager/mod.rs b/crates/vm/src/system/memory/manager/mod.rs index 595257a3c4..21ceb931a9 100644 --- a/crates/vm/src/system/memory/manager/mod.rs +++ b/crates/vm/src/system/memory/manager/mod.rs @@ -67,22 +67,22 @@ pub struct TimestampedValues { pub type MemoryControllerRef = Rc>>; -/// A equipartition of memory, with timestamps and values. +/// An equipartition of memory, with timestamps and values. /// /// The key is a pair `(address_space, label)`, where `label` is the index of the block in the /// partition. I.e., the starting address of the block is `(address_space, label * N)`. /// /// If a key is not present in the map, then the block is uninitialized (and therefore zero). pub type TimestampedEquipartition = - BTreeMap<(F, usize), TimestampedValues>; + BTreeMap<(u32, u32), TimestampedValues>; -/// A equipartition of memory values. +/// An equipartition of memory values. /// /// The key is a pair `(address_space, label)`, where `label` is the index of the block in the /// partition. I.e., the starting address of the block is `(address_space, label * N)`. /// /// If a key is not present in the map, then the block is uninitialized (and therefore zero). -pub type Equipartition = BTreeMap<(F, usize), [F; N]>; +pub type Equipartition = BTreeMap<(u32, u32), [F; N]>; #[derive(Debug, Getters)] pub struct MemoryController { @@ -346,6 +346,7 @@ impl MemoryController { } pub fn read(&mut self, address_space: F, pointer: F) -> MemoryReadRecord { + let address_space_u32 = address_space.as_canonical_u32(); let ptr_u32 = pointer.as_canonical_u32(); assert!( address_space == F::ZERO || ptr_u32 < (1 << self.mem_config.pointer_max_bits), @@ -367,16 +368,14 @@ impl MemoryController { }; } - let (record, adapter_records) = self - .memory - .read::(address_space.as_canonical_u32() as usize, ptr_u32 as usize); + let (record, adapter_records) = self.memory.read::(address_space_u32, ptr_u32); for record in adapter_records { self.access_adapters.add_record(record); } for i in 0..N as u32 { - let ptr = F::from_canonical_u32(ptr_u32 + i); - self.interface_chip.touch_address(address_space, ptr); + self.interface_chip + .touch_address(address_space_u32, ptr_u32 + i); } record @@ -393,9 +392,9 @@ impl MemoryController { /// /// Any value returned is unconstrained. pub fn unsafe_read(&self, addr_space: F, ptr: F) -> [F; N] { - let addr_space = addr_space.as_canonical_u32() as usize; - let ptr = ptr.as_canonical_u32() as usize; - from_fn(|i| self.memory.get(addr_space, ptr + i)) + let addr_space = addr_space.as_canonical_u32(); + let ptr = ptr.as_canonical_u32(); + from_fn(|i| self.memory.get(addr_space, ptr + i as u32)) } pub fn write_cell(&mut self, address_space: F, pointer: F, data: F) -> MemoryWriteRecord { @@ -409,24 +408,21 @@ impl MemoryController { data: [F; N], ) -> MemoryWriteRecord { assert_ne!(address_space, F::ZERO); + let address_space_u32 = address_space.as_canonical_u32(); let ptr_u32 = pointer.as_canonical_u32(); assert!( ptr_u32 < (1 << self.mem_config.pointer_max_bits), "memory out of bounds: {ptr_u32:?}", ); - let (record, adapter_records) = self.memory.write( - address_space.as_canonical_u32() as usize, - ptr_u32 as usize, - data, - ); + let (record, adapter_records) = self.memory.write(address_space_u32, ptr_u32, data); for record in adapter_records { self.access_adapters.add_record(record); } for i in 0..N as u32 { - let ptr = F::from_canonical_u32(ptr_u32 + i); - self.interface_chip.touch_address(address_space, ptr); + self.interface_chip + .touch_address(address_space_u32, ptr_u32 + i); } record @@ -715,12 +711,11 @@ impl MemoryAuxColsFactory { pub fn memory_image_to_equipartition( memory_image: MemoryImage, -) -> Equipartition { +) -> Equipartition { let mut result = Equipartition::new(); for ((addr_space, addr), word) in memory_image { - let addr_usize = addr.as_canonical_u32() as usize; - let shift = addr_usize % N; - let key = (addr_space, addr_usize / N); + let shift = (addr % N as u32) as usize; + let key = (addr_space, addr / N as u32); result.entry(key).or_insert([F::ZERO; N])[shift] = word; } result diff --git a/crates/vm/src/system/memory/merkle/mod.rs b/crates/vm/src/system/memory/merkle/mod.rs index ad01e25c8b..6dc75ecf6c 100644 --- a/crates/vm/src/system/memory/merkle/mod.rs +++ b/crates/vm/src/system/memory/merkle/mod.rs @@ -17,7 +17,7 @@ mod tests; #[derive(Debug)] pub struct MemoryMerkleChip { pub air: MemoryMerkleAir, - touched_nodes: FxHashSet<(usize, usize, usize)>, + touched_nodes: FxHashSet<(usize, u32, u32)>, num_touched_nonleaves: usize, final_state: Option>, overridden_height: Option, @@ -56,7 +56,7 @@ impl MemoryMerkleChip { self.overridden_height = Some(override_height); } - fn touch_node(&mut self, height: usize, as_label: usize, address_label: usize) { + fn touch_node(&mut self, height: usize, as_label: u32, address_label: u32) { if self.touched_nodes.insert((height, as_label, address_label)) { assert_ne!(height, self.air.memory_dimensions.overall_height()); if height != 0 { @@ -70,11 +70,11 @@ impl MemoryMerkleChip { } } - pub fn touch_address(&mut self, address_space: F, address: F) { + pub fn touch_address(&mut self, address_space: u32, address: u32) { self.touch_node( 0, - (address_space.as_canonical_u32() as usize) - self.air.memory_dimensions.as_offset, - (address.as_canonical_u32() as usize) / CHUNK, + address_space - self.air.memory_dimensions.as_offset, + address / CHUNK as u32, ); } } diff --git a/crates/vm/src/system/memory/merkle/tests/mod.rs b/crates/vm/src/system/memory/merkle/tests/mod.rs index fb800c2d13..11353a1a0f 100644 --- a/crates/vm/src/system/memory/merkle/tests/mod.rs +++ b/crates/vm/src/system/memory/merkle/tests/mod.rs @@ -6,11 +6,8 @@ use std::{ }; use openvm_stark_backend::{ - interaction::InteractionType, - p3_field::{AbstractField, PrimeField32}, - p3_matrix::dense::RowMajorMatrix, - prover::types::AirProofInput, - Chip, ChipUsageGetter, + interaction::InteractionType, p3_field::AbstractField, p3_matrix::dense::RowMajorMatrix, + prover::types::AirProofInput, Chip, ChipUsageGetter, }; use openvm_stark_sdk::{ config::baby_bear_poseidon2::BabyBearPoseidon2Engine, @@ -40,7 +37,7 @@ const COMPRESSION_BUS: DirectCompressionBus = DirectCompressionBus(POSEIDON2_DIR fn test( memory_dimensions: MemoryDimensions, initial_memory: &Equipartition, - touched_labels: BTreeSet<(BabyBear, usize)>, + touched_labels: BTreeSet<(u32, u32)>, final_memory: &Equipartition, ) { let MemoryDimensions { @@ -52,7 +49,7 @@ fn test( // checking validity of test data for (&(address_space, label), value) in final_memory { - assert!((address_space.as_canonical_u32() as usize) - as_offset < (1 << as_height)); + assert!(address_space - as_offset < (1 << as_height)); assert!(label < (1 << address_height)); if initial_memory.get(&(address_space, label)) != Some(value) { assert!(touched_labels.contains(&(address_space, label))); @@ -75,11 +72,8 @@ fn test( let mut chip = MemoryMerkleChip::::new(memory_dimensions, merkle_bus, COMPRESSION_BUS); for &(address_space, label) in touched_labels.iter() { - for i in 0..CHUNK { - chip.touch_address( - address_space, - BabyBear::from_canonical_usize(label * CHUNK + i), - ); + for i in 0..CHUNK as u32 { + chip.touch_address(address_space, label * CHUNK as u32 + i); } } @@ -96,8 +90,8 @@ fn test( let mut interaction = |interaction_type: InteractionType, is_compress: bool, height: usize, - as_label: usize, - address_label: usize, + as_label: u32, + address_label: u32, hash: [BabyBear; CHUNK]| { let expand_direction = if is_compress { BabyBear::NEG_ONE @@ -111,8 +105,8 @@ fn test( dummy_interaction_trace_rows.extend([ expand_direction, BabyBear::from_canonical_usize(height), - BabyBear::from_canonical_usize(as_label), - BabyBear::from_canonical_usize(address_label), + BabyBear::from_canonical_u32(as_label), + BabyBear::from_canonical_u32(address_label), ]); dummy_interaction_trace_rows.extend(hash); }; @@ -121,7 +115,7 @@ fn test( let initial_values = *initial_memory .get(&(address_space, address_label)) .unwrap_or(&[BabyBear::ZERO; CHUNK]); - let as_label = address_space.as_canonical_u32() as usize - as_offset; + let as_label = address_space - as_offset; interaction( InteractionType::Send, false, @@ -163,12 +157,12 @@ fn test( fn random_test( height: usize, - max_value: usize, + max_value: u32, mut num_initial_addresses: usize, mut num_touched_addresses: usize, ) { let mut rng = create_seeded_rng(); - let mut next_usize = || rng.next_u64() as usize; + let mut next_u32 = || rng.next_u64() as u32; let mut initial_memory = Equipartition::new(); let mut final_memory = Equipartition::new(); @@ -176,15 +170,15 @@ fn random_test( let mut touched_labels = BTreeSet::new(); while num_initial_addresses != 0 || num_touched_addresses != 0 { - let address_space = BabyBear::from_canonical_usize((next_usize() & 1) + 1); - let label = next_usize() % (1 << height); + let address_space = (next_u32() & 1) + 1; + let label = next_u32() % (1 << height); if seen_labels.insert(label) { - let is_initial = next_usize() & 1 == 0; + let is_initial = next_u32() & 1 == 0; let initial_values = - array::from_fn(|_| BabyBear::from_canonical_usize(next_usize() % max_value)); - let is_touched = next_usize() & 1 == 0; - let value_changes = next_usize() & 1 == 0; + array::from_fn(|_| BabyBear::from_canonical_u32(next_u32() % max_value)); + let is_touched = next_u32() & 1 == 0; + let value_changes = next_u32() & 1 == 0; if is_initial && num_initial_addresses != 0 { num_initial_addresses -= 1; @@ -195,9 +189,8 @@ fn random_test( num_touched_addresses -= 1; touched_labels.insert((address_space, label)); if value_changes || !is_initial { - let changed_values = array::from_fn(|_| { - BabyBear::from_canonical_usize(next_usize() % max_value) - }); + let changed_values = + array::from_fn(|_| BabyBear::from_canonical_u32(next_u32() % max_value)); final_memory.insert((address_space, label), changed_values); } } diff --git a/crates/vm/src/system/memory/merkle/trace.rs b/crates/vm/src/system/memory/merkle/trace.rs index 8045e7fc92..68f5caa30a 100644 --- a/crates/vm/src/system/memory/merkle/trace.rs +++ b/crates/vm/src/system/memory/merkle/trace.rs @@ -119,7 +119,7 @@ impl ChipUsageGetter for MemoryMerkleChip { memory_dimensions: MemoryDimensions, final_memory: &'a Equipartition, - touched_nodes: &'a FxHashSet<(usize, usize, usize)>, + touched_nodes: &'a FxHashSet<(usize, u32, u32)>, trace_rows: &'a mut Vec>, } @@ -128,13 +128,12 @@ impl TreeHelper<'_, CHUNK, F> { &mut self, height: usize, initial_node: &MemoryNode, - as_label: usize, - address_label: usize, + as_label: u32, + address_label: u32, hasher: &mut impl HasherChip, ) -> MemoryNode { if height == 0 { - let address_space = - F::from_canonical_usize(as_label + self.memory_dimensions.as_offset); + let address_space = as_label + self.memory_dimensions.as_offset; let leaf_values = *self .final_memory .get(&(address_space, address_label)) @@ -215,8 +214,8 @@ impl TreeHelper<'_, CHUNK, F> { fn add_trace_row( &mut self, parent_height: usize, - as_label: usize, - address_label: usize, + as_label: u32, + address_label: u32, node: &MemoryNode, direction_changes: Option<[bool; 2]>, ) { @@ -232,8 +231,8 @@ impl TreeHelper<'_, CHUNK, F> { height_section: F::from_bool(parent_height > self.memory_dimensions.address_height), parent_height: F::from_canonical_usize(parent_height), is_root: F::from_bool(parent_height == self.memory_dimensions.overall_height()), - parent_as_label: F::from_canonical_usize(as_label), - parent_address_label: F::from_canonical_usize(address_label), + parent_as_label: F::from_canonical_u32(as_label), + parent_address_label: F::from_canonical_u32(address_label), parent_hash: *hash, left_child_hash: left.hash(), right_child_hash: right.hash(), diff --git a/crates/vm/src/system/memory/persistent.rs b/crates/vm/src/system/memory/persistent.rs index 57eb71d633..bbd0ab0b9e 100644 --- a/crates/vm/src/system/memory/persistent.rs +++ b/crates/vm/src/system/memory/persistent.rs @@ -87,7 +87,7 @@ impl Air for PersistentBoundaryA // direction = -1 => is_final = 1 local.expand_direction.into(), AB::Expr::ZERO, - local.address_space - AB::F::from_canonical_usize(self.memory_dims.as_offset), + local.address_space - AB::F::from_canonical_u32(self.memory_dims.as_offset), local.leaf_label.into(), ]; expand_fields.extend(local.hash.map(Into::into)); @@ -128,14 +128,14 @@ pub struct PersistentBoundaryChip { #[derive(Debug)] enum TouchedLabels { - Running(FxHashSet<(F, usize)>), + Running(FxHashSet<(u32, u32)>), Final(Vec>), } #[derive(Debug)] struct FinalTouchedLabel { - address_space: F, - label: usize, + address_space: u32, + label: u32, init_values: [F; CHUNK], final_values: [F; CHUNK], init_exists: bool, @@ -151,7 +151,7 @@ impl Default for TouchedLabels { } impl TouchedLabels { - fn touch(&mut self, address_space: F, label: usize) { + fn touch(&mut self, address_space: u32, label: u32) { match self { TouchedLabels::Running(touched_labels) => { touched_labels.insert((address_space, label)); @@ -190,8 +190,8 @@ impl PersistentBoundaryChip { self.overridden_height = Some(overridden_height); } - pub fn touch_address(&mut self, address_space: F, pointer: F) { - let label = pointer.as_canonical_u32() as usize / CHUNK; + pub fn touch_address(&mut self, address_space: u32, pointer: u32) { + let label = pointer / CHUNK as u32; self.touched_labels.touch(address_space, label); } @@ -272,8 +272,8 @@ where let (initial_row, final_row) = row.split_at_mut(width); *initial_row.borrow_mut() = PersistentBoundaryCols { expand_direction: Val::::ONE, - address_space: touched_label.address_space, - leaf_label: Val::::from_canonical_usize(touched_label.label), + address_space: Val::::from_canonical_u32(touched_label.address_space), + leaf_label: Val::::from_canonical_u32(touched_label.label), values: touched_label.init_values, hash: touched_label.init_hash, timestamp: if touched_label.init_exists { @@ -285,8 +285,8 @@ where *final_row.borrow_mut() = PersistentBoundaryCols { expand_direction: Val::::NEG_ONE, - address_space: touched_label.address_space, - leaf_label: Val::::from_canonical_usize(touched_label.label), + address_space: Val::::from_canonical_u32(touched_label.address_space), + leaf_label: Val::::from_canonical_u32(touched_label.label), values: touched_label.final_values, hash: touched_label.final_hash, timestamp: Val::::from_canonical_u32(touched_label.final_timestamp), diff --git a/crates/vm/src/system/memory/tree/mod.rs b/crates/vm/src/system/memory/tree/mod.rs index 00a8281152..51e58a2ab8 100644 --- a/crates/vm/src/system/memory/tree/mod.rs +++ b/crates/vm/src/system/memory/tree/mod.rs @@ -66,9 +66,9 @@ impl MemoryNode { } fn from_memory( - memory: &BTreeMap, + memory: &BTreeMap, height: usize, - from: usize, + from: u64, hasher: &impl Hasher, ) -> MemoryNode { let mut range = memory.range(from..from + (1 << height)); diff --git a/crates/vm/src/system/memory/tree/public_values.rs b/crates/vm/src/system/memory/tree/public_values.rs index 385e1582bd..f31ca63226 100644 --- a/crates/vm/src/system/memory/tree/public_values.rs +++ b/crates/vm/src/system/memory/tree/public_values.rs @@ -8,7 +8,7 @@ use crate::{ system::memory::{dimensions::MemoryDimensions, tree::MemoryNode, Equipartition}, }; -pub const PUBLIC_VALUES_ADDRESS_SPACE_OFFSET: usize = 2; +pub const PUBLIC_VALUES_ADDRESS_SPACE_OFFSET: u32 = 2; /// Merkle proof for user public values in the memory state. #[derive(Clone, Debug, Serialize, Deserialize)] @@ -111,11 +111,8 @@ pub fn extract_public_values( final_memory: &Equipartition, ) -> Vec { // All (addr, value) pairs in the public value address space. - let f_as_start = - F::from_canonical_usize(PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset); - let f_as_end = F::from_canonical_usize( - PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset + 1, - ); + let f_as_start = PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset; + let f_as_end = PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset + 1; let used_pvs: Vec<_> = final_memory .range((f_as_start, 0)..(f_as_end, 0)) @@ -123,7 +120,7 @@ pub fn extract_public_values( value .iter() .enumerate() - .map(|(i, &v)| (*block_id * CHUNK + i, v)) + .map(|(i, &v)| (*block_id as usize * CHUNK + i, v)) }) .collect(); if let Some(last_pv) = used_pvs.last() { @@ -161,13 +158,9 @@ mod tests { vm_config.memory_config.as_height = 4; vm_config.memory_config.pointer_max_bits = 5; let memory_dimensions = vm_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 num_public_values = 16; - let memory: MemoryImage = [((pv_as, F::from_canonical_u32(15)), F::ONE)] - .into_iter() - .collect(); + let memory: MemoryImage = [((pv_as, 15), F::ONE)].into_iter().collect(); let mut expected_pvs = F::zero_vec(num_public_values); expected_pvs[15] = F::ONE; diff --git a/crates/vm/src/system/memory/volatile/mod.rs b/crates/vm/src/system/memory/volatile/mod.rs index 04e0c68579..a831d768b3 100644 --- a/crates/vm/src/system/memory/volatile/mod.rs +++ b/crates/vm/src/system/memory/volatile/mod.rs @@ -1,6 +1,5 @@ use std::{ borrow::{Borrow, BorrowMut}, - collections::HashSet, sync::Arc, }; @@ -24,6 +23,7 @@ use openvm_stark_backend::{ rap::{AnyRap, BaseAirWithPublicValues, PartitionedBaseAir}, Chip, ChipUsageGetter, }; +use rustc_hash::FxHashSet; use super::TimestampedEquipartition; use crate::system::memory::{ @@ -133,7 +133,7 @@ impl Air for VolatileBoundaryAir { #[derive(Debug)] pub struct VolatileBoundaryChip { pub air: VolatileBoundaryAir, - touched_addresses: HashSet<(F, F)>, + touched_addresses: FxHashSet<(u32, u32)>, range_checker: Arc, overridden_height: Option, final_memory: Option>, @@ -154,18 +154,18 @@ impl VolatileBoundaryChip { pointer_max_bits, range_bus, ), - touched_addresses: HashSet::new(), + touched_addresses: FxHashSet::default(), range_checker, overridden_height: None, final_memory: None, } } - pub fn touch_address(&mut self, addr_space: F, pointer: F) { + pub fn touch_address(&mut self, addr_space: u32, pointer: u32) { self.touched_addresses.insert((addr_space, pointer)); } - pub fn all_addresses(&self) -> Vec<(F, F)> { + pub fn all_addresses(&self) -> Vec<(u32, u32)> { self.touched_addresses.iter().cloned().collect() } } @@ -220,8 +220,8 @@ where // `pointer` is the same as `label` since the equipartition has block size 1 let [data] = timestamped_values.values; let row: &mut VolatileBoundaryCols<_> = row.borrow_mut(); - row.addr_space = *addr_space; - row.pointer = Val::::from_canonical_usize(*ptr); + row.addr_space = Val::::from_canonical_u32(*addr_space); + row.pointer = Val::::from_canonical_u32(*ptr); row.initial_data = Val::::ZERO; row.final_data = data; row.final_timestamp = Val::::from_canonical_u32(timestamped_values.timestamp); @@ -235,7 +235,10 @@ where ( &self.range_checker, &[row.addr_space, row.pointer], - &[next_addr_space, Val::::from_canonical_usize(next_ptr)], + &[ + Val::::from_canonical_u32(next_addr_space), + Val::::from_canonical_u32(next_ptr), + ], ), ((&mut row.addr_lt_aux).into(), &mut out), ); diff --git a/crates/vm/src/system/memory/volatile/tests.rs b/crates/vm/src/system/memory/volatile/tests.rs index eaaf3bb674..cec8f05cd9 100644 --- a/crates/vm/src/system/memory/volatile/tests.rs +++ b/crates/vm/src/system/memory/volatile/tests.rs @@ -2,10 +2,7 @@ use std::{collections::HashSet, iter, sync::Arc}; use openvm_circuit_primitives::var_range::{VariableRangeCheckerBus, VariableRangeCheckerChip}; use openvm_stark_backend::{ - p3_field::{AbstractField, PrimeField32}, - p3_matrix::dense::RowMajorMatrix, - prover::types::AirProofInput, - Chip, + p3_field::AbstractField, p3_matrix::dense::RowMajorMatrix, prover::types::AirProofInput, Chip, }; use openvm_stark_sdk::{ config::baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine}, @@ -30,17 +27,17 @@ fn boundary_air_test() { const MEMORY_BUS: usize = 1; const RANGE_CHECKER_BUS: usize = 3; - const MAX_ADDRESS_SPACE: usize = 4; + const MAX_ADDRESS_SPACE: u32 = 4; const LIMB_BITS: usize = 15; - const MAX_VAL: usize = 1 << LIMB_BITS; + const MAX_VAL: u32 = 1 << LIMB_BITS; const DECOMP: usize = 8; let memory_bus = MemoryBus(MEMORY_BUS); let num_addresses = 10; let mut distinct_addresses = HashSet::new(); while distinct_addresses.len() < num_addresses { - let addr_space = Val::from_canonical_usize(rng.gen_range(0..MAX_ADDRESS_SPACE)); - let pointer = Val::from_canonical_usize(rng.gen_range(0..MAX_VAL)); + let addr_space = rng.gen_range(0..MAX_ADDRESS_SPACE); + let pointer = rng.gen_range(0..MAX_VAL); distinct_addresses.insert((addr_space, pointer)); } @@ -52,11 +49,11 @@ fn boundary_air_test() { let mut final_memory = TimestampedEquipartition::new(); for (addr_space, pointer) in distinct_addresses.iter().cloned() { - let final_data = Val::from_canonical_usize(rng.gen_range(0..MAX_VAL)); + let final_data = Val::from_canonical_u32(rng.gen_range(0..MAX_VAL)); let final_clk = rng.gen_range(1..MAX_VAL) as u32; final_memory.insert( - (addr_space, pointer.as_canonical_u32() as usize), + (addr_space, pointer), TimestampedValues { values: [final_data], timestamp: final_clk, @@ -75,8 +72,8 @@ fn boundary_air_test() { .flat_map(|(addr_space, pointer)| { vec![ Val::ONE, - *addr_space, - *pointer, + Val::from_canonical_u32(*addr_space), + Val::from_canonical_u32(*pointer), Val::ZERO, Val::ZERO, Val::ONE, @@ -91,14 +88,12 @@ fn boundary_air_test() { distinct_addresses .iter() .flat_map(|(addr_space, pointer)| { - let timestamped_value = final_memory - .get(&(*addr_space, pointer.as_canonical_u32() as usize)) - .unwrap(); + let timestamped_value = final_memory.get(&(*addr_space, *pointer)).unwrap(); vec![ Val::ONE, - *addr_space, - *pointer, + Val::from_canonical_u32(*addr_space), + Val::from_canonical_u32(*pointer), timestamped_value.values[0], Val::from_canonical_u32(timestamped_value.timestamp), Val::ONE, diff --git a/crates/vm/tests/integration_test.rs b/crates/vm/tests/integration_test.rs index c07790c4e4..224388ac2a 100644 --- a/crates/vm/tests/integration_test.rs +++ b/crates/vm/tests/integration_test.rs @@ -352,12 +352,9 @@ fn test_vm_initial_memory() { ), ]); - let init_memory: BTreeMap<_, _> = [( - (BabyBear::ONE, BabyBear::from_canonical_u32(7)), - BabyBear::from_canonical_u32(101), - )] - .into_iter() - .collect(); + let init_memory: BTreeMap<_, _> = [((1, 7), BabyBear::from_canonical_u32(101))] + .into_iter() + .collect(); let config = NativeConfig::aggregation(0, 3).with_continuations(); let exe = VmExe {