From f95bb70c198b7c173c883f6b3d29c9c334940606 Mon Sep 17 00:00:00 2001 From: Galin Chung Nguyen Date: Tue, 11 Feb 2025 12:40:41 +0700 Subject: [PATCH 1/5] Update test to enhance code coverage --- zkmemory/Cargo.toml | 4 ++ zkmemory/README.md | 85 +++++++++++++++++++++++++---------------- zkmemory/src/config.rs | 52 ++++++++++++++++++++----- zkmemory/src/machine.rs | 85 +++++++++++++++++++++++++++++++++++++---- 4 files changed, 176 insertions(+), 50 deletions(-) diff --git a/zkmemory/Cargo.toml b/zkmemory/Cargo.toml index 0a1feae..4d92154 100644 --- a/zkmemory/Cargo.toml +++ b/zkmemory/Cargo.toml @@ -32,3 +32,7 @@ nova-snark = { workspace = true } bellpepper-core = { workspace = true } poseidon = { path = "../poseidon" } arecibo = { workspace = true } + +[features] +default = ["no_std"] +no_std = [] \ No newline at end of file diff --git a/zkmemory/README.md b/zkmemory/README.md index aeded9e..e54fe0c 100644 --- a/zkmemory/README.md +++ b/zkmemory/README.md @@ -15,42 +15,61 @@ $ cargo install cargo-llvm-cov $ cargo llvm-cov --html --open ``` -Right now, our code coverage is `80%`: +Right now, our code line coverage is `96.10%`: ```text -running 17 tests -test tests::sm256_read_empty_cell ... ok -test tests::sm256_write_one_cell_read_two_cell ... ok -test tests::sm256_write_read_one_cell ... ok -test tests::sm256_write_two_cell_read_one_cell ... ok -test tests::sm32_read_empty_cell ... ok -test tests::sm32_read_prohibited_cell - should panic ... ok -test tests::sm32_write_one_cell_read_two_cells ... ok -test tests::sm32_write_read_one_cell ... ok -test tests::sm32_write_two_cells_read_one_cells ... ok -test tests::u256_arithmetic_test ... ok -test tests::u256_test ... ok -test tests::u32_register_functional ... ok -test tests::u32_stack_functional ... ok -test tests::u32_stack_overflow - should panic ... ok -test tests::u32_stack_underflow - should panic ... ok -test tests::u32_test ... ok -test tests::u64_test ... ok - -test result: ok. 17 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s - -Filename Regions Missed Regions Cover Functions Missed Functions Executed Lines Missed Lines Cover Branches Missed Branches Cover ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ -base.rs 18 0 100.00% 18 0 100.00% 54 0 100.00% 0 0 - -config.rs 14 5 64.29% 9 4 55.56% 55 22 60.00% 0 0 - -error.rs 13 11 15.38% 4 3 25.00% 13 12 7.69% 0 0 - -lib.rs 87 1 98.85% 34 0 100.00% 172 1 99.42% 0 0 - -machine.rs 74 24 67.57% 25 8 68.00% 144 25 82.64% 0 0 - -memory.rs 18 2 88.89% 9 2 77.78% 91 4 95.60% 0 0 - ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ -TOTAL 224 43 80.80% 99 17 82.83% 529 64 87.90% 0 0 - +running 86 tests +test constraints::original_memory_circuit::tests::also_test_invalid_time_order - should panic ... ok +test constraints::original_memory_circuit::tests::test_identical_trace - should panic ... ok +test constraints::original_memory_circuit::tests::test_invalid_time_order - should panic ... ok +test constraints::original_memory_circuit::tests::test_multiple_traces ... ok +test constraints::original_memory_circuit::tests::test_one_trace ... ok +test constraints::original_memory_circuit::tests::test_wrong_starting_time - should panic ... ok +test constraints::permutation_circuit::tests::check_permutation_with_trace_records ... ok +test constraints::permutation_circuit::tests::check_trace_record_mapping ... ok +test constraints::permutation_circuit::tests::check_wrong_permutation - should panic ... ok +test constraints::permutation_circuit::tests::test_functionality ... ok +test constraints::permutation_circuit::tests::test_inequal_lengths - should panic ... ok +test constraints::sorted_memory_circuit::test::equal_address_and_time_log - should panic ... ok +test constraints::sorted_memory_circuit::test::invalid_read - should panic ... ok +test constraints::sorted_memory_circuit::test::invalid_read2 - should panic ... ok +test constraints::sorted_memory_circuit::test::invalid_read3 - should panic ... ok +test constraints::sorted_memory_circuit::test::non_first_write_access_for_two_traces - should panic ... ok +test constraints::sorted_memory_circuit::test::test_error_invalid_instruction - should panic ... ok +... + +test result: ok. 85 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 6738.79s + ``` +| Filename | Function Coverage | Line Coverage | Region Coverage | +| :---: | :---: | :---: | :---: | +| poseidon/src/circuit.rs | 78.79 \%(52/66) | 82.55 \%(388/470) | 79.33 \%(165/208)| +| poseidon/src/gadgets.rs | 80.70 \%(46/57) | 91.54 \%(292/319) | 81.10 \%(133/164)| +| poseidon/src/poseidon hash.rs | 92.00 \%(23/25) | 87.68 \%(178/203) | 92.05 \%(81/88) | +| zkmemory/src/base.rs | 100.00% (30/30) | 100.00% (94/94) | 100.00% (30/30) | +| zkmemory/src/commitment/extends.rs | 100.00% (2/2) | 100.00% (17/17) | 100.00% (6/6) | +| zkmemory/src/commitment/kzg.rs | 95.83% (23/24) | 99.45% (364/366) | 94.52% (69/73) | +| zkmemory/src/commitment/merkle tree.rs | 84.06% (58/69) | 96.33% (473/491) | 84.83% (151/178) | +| zkmemory/src/commitment/verkle tree.rs | 85.71% (48/56) | 96.37% (478/496) | 84.93% (124/146) | +| zkmemory/src/config.rs | 100.00% (10/10) | 100.00% (117/117) | 100.00% (31/31) | +| zkmemory/src/constraints/consistency check circuit.rs | 80.00% (4/5) | 96.97% (96/99) | 82.61% (19/23) | +| zkmemory/src/constraints/gadgets.rs | 95.24% (40/42) | 99.40% (331/333) | 95.41% (104/109) | +| zkmemory/src/constraints/helper.rs | 100.00% (13/13) | 99.70% (336/337) | 94.74% (18/19) | +| zkmemory/src/constraints/original memory circuit. rs | 69.23% (27/39) | 95.87% (325/339) | 76.56 \%(98/128) | +| zkmemory/src/constraints/permutation circuit.rs | 83.33% (35/42) | 97.09% (334/344) | 84.85% (112/132) | +| zkmemory/src/constraints/sorted memory circuit.rs | 75.81% (47/62) | 97.01% (551/568) | 80.54% (149/185) | +| zkmemory/src/error.rs | 100.00% (2/2) | 100.00% (34/34) | 100.00% (17/17) | +| zkmemory/src/lib.rs | 100.00% (4/4) | 100.00% (119/119) | 100.00% (47/47) | +| zkmemory/src/machine.rs | 100.00% (48/48) | 96.98% (545/562) | 90.30% (149/165) | +| zkmemory/src/nova/memory consistency circuit.rs | 71.05% (27/38) | 94.95% (207/218) | 89.11% (90/101) | +| zkmemory/src/nova/poseidon parameters.rs | 100.00% (4/4) | 100.00% (16/16) | 100.00% (4/4) | +| zkmemory/src/nova/testcases.rs | 100.00% (12/12) | 100.00% (513/513) | 100.00% (37/37) | +| zkmemory/src/supernova/memory consistency circuit.rs | 75.00% (42/56) | 95.79% (341/356) | 89.80% (132/147) | +| zkmemory/src/supernova/poseidon_parameter.rs | 100.00% (4/4) | 100.00% (16/16) | 100.00% (4/4) | +| zkmemory/src/supernova/testcases.rs | 100.00% (7/7) | 99.69% (318/319) | 97.62% (41/42) | +| Totals | 84.91% (608/716) | 96.10% (6483/6746) | 86.90% (1811/2084) | + ## Overview The idea is to create an independent module that can be used by any zkVM. You might aware that the memory can be constructed as a simple state machine with `2` instructions `READ` and `WRITE`, and configurable `WORD_SIZE`. Our memory state machine is only able access the exactly `WORD_SIZE` for every executed instruction. That is, if you want to access arbitrary data size, it must be translated to multiple accesses. @@ -61,7 +80,7 @@ These instructions need to be satisfied following conditions: - `READ` on a memory was not wrote should return `0` - Every`READ` access for the same location, must have the value to be equal to the previous `WRITE`. - **`WRITE` instruction** - - Every `WRITE` access must write on writable memory chunks _(some areas of the memory might be read only)_. +- Every `WRITE` access must write on writable memory chunks _(some areas of the memory might be read only)_. ## Features diff --git a/zkmemory/src/config.rs b/zkmemory/src/config.rs index 8979fd1..2d1cc05 100644 --- a/zkmemory/src/config.rs +++ b/zkmemory/src/config.rs @@ -1,5 +1,6 @@ use crate::base::Base; use crate::machine::Register; +extern crate std; /// Memory section #[derive(Debug, Clone, Copy)] @@ -51,7 +52,7 @@ pub struct ConfigArgs { pub stack_depth: T, /// Number of registers pub no_register: T, - /// Buffer size + /// Buffer size in words pub buffer_size: T, } @@ -77,11 +78,12 @@ where /// Create a new config for given arguments pub fn new(word_size: T, args: ConfigArgs) -> Self { if args.head_layout { + // [stack - buffer - register - buffer - memory] let stack_lo = T::MIN; let stack_hi = stack_lo + (args.stack_depth * word_size); - let register_lo = stack_hi + args.buffer_size; + let register_lo = stack_hi + args.buffer_size * word_size; let register_hi = register_lo + (args.no_register * word_size); - let memory_lo = register_hi + args.buffer_size; + let memory_lo = register_hi + args.buffer_size * word_size; let memory_hi = T::MAX; Self { word_size, @@ -92,17 +94,18 @@ where memory: AllocatedSection(memory_lo, memory_hi), } } else { - let length = + // [memory - buffer - stack - buffer - register] + let stack_register_buffer_total_size = (args.stack_depth + args.no_register + args.buffer_size + args.buffer_size) * word_size; - let stack_lo = T::MAX - length; + let stack_lo = T::MAX - stack_register_buffer_total_size; let remain = stack_lo % word_size; - let stack_lo = stack_lo - remain + word_size; + let stack_lo = stack_lo - remain; // Align to the nearest previous word-aligned address to ensure sufficient allocation for stack and register sections. let stack_hi = stack_lo + (args.stack_depth * word_size); - let register_lo = stack_hi + args.buffer_size; + let register_lo = stack_hi + args.buffer_size * word_size; let register_hi = register_lo + (args.no_register * word_size); let memory_lo = T::MIN; - let memory_hi = T::MAX - length; + let memory_hi = stack_lo - args.buffer_size * word_size; Self { word_size, @@ -122,6 +125,11 @@ where self.register.low() + (T::from(index) * self.word_size), ) } + + /// Get the buffer size in bytes + pub fn buffer_size_in_bytes(&self) -> T { + self.buffer_size * self.word_size + } } #[cfg(test)] @@ -129,6 +137,7 @@ mod tests { use super::ConfigArgs; use crate::base::{Base, B256}; use crate::config::{Config, DefaultConfig}; + extern crate std; impl PartialEq for ConfigArgs { fn eq(&self, other: &Self) -> bool { @@ -156,17 +165,42 @@ mod tests { let config = Config::::new(B256::from(32), DefaultConfig::default_config()); assert!(config.memory.contain(B256::MAX - B256::from(1))); + assert_eq!(config.stack.low(), B256::from(0)); + assert_eq!( + config.register.low(), + B256::from(config.stack.high() + config.buffer_size_in_bytes()) + ); + assert_eq!( + config.memory.low(), + B256::from(config.register.high() + config.buffer_size_in_bytes()) + ); + assert_eq!(config.memory.high(), B256::MAX); + + let no_register = B256::from(32); // Test tail layout let config = Config::::new( B256::from(32), ConfigArgs { head_layout: false, stack_depth: B256::from(1024), - no_register: B256::from(32), + no_register, buffer_size: B256::from(32), }, ); assert!(config.memory.contain(B256::from(0x10000f))); + assert_eq!(config.memory.low(), B256::from(0)); + assert_eq!( + config.memory.high(), + B256::from(config.stack.low() - config.buffer_size_in_bytes()) + ); + assert_eq!( + config.stack.high(), + B256::from(config.register.low() - config.buffer_size_in_bytes()) + ); + assert_eq!( + config.register.high(), + config.register.low() + no_register * config.word_size + ); // Test register config.create_register(0); diff --git a/zkmemory/src/machine.rs b/zkmemory/src/machine.rs index 729078a..15b34eb 100644 --- a/zkmemory/src/machine.rs +++ b/zkmemory/src/machine.rs @@ -1,4 +1,5 @@ extern crate alloc; +extern crate std; use crate::{base::Base, error::Error}; use alloc::vec::Vec; use rbtree::RBTree; @@ -264,7 +265,7 @@ where if remain.is_zero() { let time_log = self.ro_context().time_log(); // Write on a cell - self.context().memory().insert(address, value); + self.context().memory().replace_or_insert(address, value); self.track(Self::TraceRecord::new( time_log, self.ro_context().stack_depth(), @@ -581,13 +582,14 @@ mod tests { error::Error, machine::{ AbstractContext, AbstractInstruction, AbstractMachine, AbstractMemoryMachine, - AbstractRegisterMachine, AbstractStackMachine, CellInteraction, Register, TraceRecord, + AbstractRegisterMachine, AbstractStackMachine, AbstractTraceRecord, CellInteraction, + MemoryInstruction, Register, TraceRecord, }, }; extern crate alloc; extern crate std; use alloc::{vec, vec::Vec}; - use core::marker::PhantomData; + use core::{cmp::Ordering, marker::PhantomData}; use rbtree::RBTree; /// My instruction set for the machine @@ -743,11 +745,18 @@ mod tests { } MyInstruction::Swap(reg) => { match machine.pop().expect("Unable to pop value from stack") { - (_, CellInteraction::SingleCell(_op, _addr, value)) => { - machine - .push(value) - .expect("Unable to push register's value to stack"); - machine.set(*reg, value).expect("Unable to set register"); + (_, CellInteraction::SingleCell(_op, _addr, stack_top_value)) => { + let register_cell = + machine.get(*reg).expect("Unable to access register"); + if let CellInteraction::SingleCell(_, _, register_value) = register_cell + { + machine + .push(register_value) + .expect("Unable to push register's value to stack"); + machine + .set(*reg, stack_top_value) + .expect("Unable to set register"); + } } _ => panic!("Stack unable to be two cells"), }; @@ -927,6 +936,8 @@ mod tests { for instruction in program { sm.exec(&instruction); } + assert_eq!(sm.get_memory_address().0, B256::from(35840)); + assert_eq!(write_chunk, sm.dummy_read(base + B256::from(32))); assert_eq!(write_chunk, sm.dummy_read(base + B256::from(32))); } @@ -1012,6 +1023,18 @@ mod tests { sm.exec(&instruction); } + let first_trace_record = *sm.trace().first().unwrap(); + let last_trace_record = *sm.trace().last().unwrap(); + assert_eq!(last_trace_record.stack_depth(), 0); + assert_eq!(last_trace_record.value(), B256::from(170)); + assert_eq!(last_trace_record.address, base + B256::from(192)); + assert_eq!(last_trace_record.instruction(), MemoryInstruction::Write); + assert_eq!( + last_trace_record.partial_cmp(&first_trace_record), + Some(Ordering::Greater) + ); + assert_eq!(sm.get_stack_depth(), 0); + assert_eq!(sm.dummy_read(base + B256::from(128)), B256::from(170)); assert_eq!(sm.dummy_read(base + B256::from(160)), B256::from(1000)); assert_eq!(sm.dummy_read(base + B256::from(192)), B256::from(170)); @@ -1056,4 +1079,50 @@ mod tests { sm.exec(&instruction); } } + + #[test] + fn test_register_machine() { + let mut sm = StateMachine::::new(DefaultConfig::default_config()); // by default there are 32 registers allocated + + assert_eq!(sm.register_start(), B256::from(33792)); + assert_eq!(sm.base_address(), B256::from(35840)); + assert_eq!(sm.stack_allocated.low(), B256::zero()); // 256 * 4 = 1024, buffer size = 32 + assert_eq!(sm.word_size(), B256::from(32)); + + let r5 = sm.new_register(5).expect("Failed to create register"); + + assert_eq!(r5.index(), 5); + sm.set(sm.r1, B256::from(70)) + .expect("Failed to set register"); + assert_eq!(sm.dummy_read(sm.r1.address()), B256::from(70)); + assert_eq!(sm.dummy_read(r5.address()), B256::from(0)); + + let program = vec![ + Instruction::Push(B256::from(1000)), + Instruction::Swap(sm.r1), + Instruction::Swap(r5), + ]; + // Execute program1 + for instruction in program { + sm.exec(&instruction); + } + + let r5_cell_interaction = sm.get(r5).expect("Failed to get register"); + if let CellInteraction::SingleCell(MemoryInstruction::Read, address, value) = + r5_cell_interaction + { + assert_eq!((address, value), (r5.address(), B256::from(70))); + } else { + panic!("Unexpected CellInteraction variant"); + } + + let r1_cell_interaction = sm.get(sm.r1).expect("Failed to get register"); + if let CellInteraction::SingleCell(MemoryInstruction::Read, address, value) = + r1_cell_interaction + { + assert_eq!((address, value), (sm.r1.address(), B256::from(1000))); + } else { + panic!("Unexpected CellInteraction variant"); + } + } } From 8ede785ab4076a090d44f66a0ed1976a2afb0188 Mon Sep 17 00:00:00 2001 From: Galin Chung Nguyen Date: Fri, 14 Feb 2025 16:55:52 +0700 Subject: [PATCH 2/5] Fix minor code style bugs --- libecvrf/examples/vrf.rs | 2 +- zkmemory/src/config.rs | 16 +++++----------- zkmemory/src/machine.rs | 4 +--- 3 files changed, 7 insertions(+), 15 deletions(-) diff --git a/libecvrf/examples/vrf.rs b/libecvrf/examples/vrf.rs index a8de03b..6ff3429 100644 --- a/libecvrf/examples/vrf.rs +++ b/libecvrf/examples/vrf.rs @@ -1,5 +1,5 @@ use libecvrf::{ - extends::{AffineExtend, ScalarExtend}, + extends::ScalarExtend, helper::{calculate_witness_address, get_address}, secp256k1::{curve::Scalar, SecretKey}, util::thread_rng, diff --git a/zkmemory/src/config.rs b/zkmemory/src/config.rs index 2d1cc05..28288e7 100644 --- a/zkmemory/src/config.rs +++ b/zkmemory/src/config.rs @@ -1,6 +1,5 @@ use crate::base::Base; use crate::machine::Register; -extern crate std; /// Memory section #[derive(Debug, Clone, Copy)] @@ -125,11 +124,6 @@ where self.register.low() + (T::from(index) * self.word_size), ) } - - /// Get the buffer size in bytes - pub fn buffer_size_in_bytes(&self) -> T { - self.buffer_size * self.word_size - } } #[cfg(test)] @@ -137,7 +131,6 @@ mod tests { use super::ConfigArgs; use crate::base::{Base, B256}; use crate::config::{Config, DefaultConfig}; - extern crate std; impl PartialEq for ConfigArgs { fn eq(&self, other: &Self) -> bool { @@ -168,11 +161,11 @@ mod tests { assert_eq!(config.stack.low(), B256::from(0)); assert_eq!( config.register.low(), - B256::from(config.stack.high() + config.buffer_size_in_bytes()) + B256::from(config.stack.high() + config.buffer_size * config.word_size) ); assert_eq!( config.memory.low(), - B256::from(config.register.high() + config.buffer_size_in_bytes()) + B256::from(config.register.high() + config.buffer_size * config.word_size) ); assert_eq!(config.memory.high(), B256::MAX); @@ -187,15 +180,16 @@ mod tests { buffer_size: B256::from(32), }, ); + assert!(config.memory.contain(B256::from(0x10000f))); assert_eq!(config.memory.low(), B256::from(0)); assert_eq!( config.memory.high(), - B256::from(config.stack.low() - config.buffer_size_in_bytes()) + B256::from(config.stack.low() - config.buffer_size * config.word_size) ); assert_eq!( config.stack.high(), - B256::from(config.register.low() - config.buffer_size_in_bytes()) + B256::from(config.register.low() - config.buffer_size * config.word_size) ); assert_eq!( config.register.high(), diff --git a/zkmemory/src/machine.rs b/zkmemory/src/machine.rs index 15b34eb..6ca80a5 100644 --- a/zkmemory/src/machine.rs +++ b/zkmemory/src/machine.rs @@ -1,5 +1,4 @@ extern crate alloc; -extern crate std; use crate::{base::Base, error::Error}; use alloc::vec::Vec; use rbtree::RBTree; @@ -587,7 +586,6 @@ mod tests { }, }; extern crate alloc; - extern crate std; use alloc::{vec, vec::Vec}; use core::{cmp::Ordering, marker::PhantomData}; use rbtree::RBTree; @@ -936,7 +934,7 @@ mod tests { for instruction in program { sm.exec(&instruction); } - assert_eq!(sm.get_memory_address().0, B256::from(35840)); + assert_eq!(sm.get_memory_address(), (B256::from(35840), B256::MAX)); assert_eq!(write_chunk, sm.dummy_read(base + B256::from(32))); assert_eq!(write_chunk, sm.dummy_read(base + B256::from(32))); } From f8819663593eefbe1882cd9f9ce12e4b842cce63 Mon Sep 17 00:00:00 2001 From: Chiro Hiro Date: Mon, 17 Feb 2025 17:01:46 +0700 Subject: [PATCH 3/5] Update memory allocation section --- zkmemory/examples/256bits-machine.rs | 12 +- zkmemory/examples/memory-consistency.rs | 12 +- zkmemory/src/config.rs | 181 ++++++++++++++---------- zkmemory/src/machine.rs | 24 ++-- 4 files changed, 127 insertions(+), 102 deletions(-) diff --git a/zkmemory/examples/256bits-machine.rs b/zkmemory/examples/256bits-machine.rs index 40da4fd..f311c39 100644 --- a/zkmemory/examples/256bits-machine.rs +++ b/zkmemory/examples/256bits-machine.rs @@ -2,7 +2,7 @@ use rbtree::RBTree; use std::{marker::PhantomData, println}; use zkmemory::{ base::{Base, B256}, - config::{AllocatedSection, Config, ConfigArgs, DefaultConfig}, + config::{Config, ConfigArgs, Section}, error::Error, impl_register_machine, impl_stack_machine, impl_state_machine, machine::{ @@ -52,18 +52,18 @@ where { // Memory memory: RBTree, - memory_allocated: AllocatedSection, + memory_allocated: Section, word_size: K, time_log: u64, // Stack - stack_allocated: AllocatedSection, + stack_allocated: Section, max_stack_depth: u64, stack_depth: u64, stack_ptr: K, // Register - register_allocated: AllocatedSection, + register_allocated: Section, /// Register r0 pub r0: Register, @@ -216,7 +216,7 @@ where V: Base, { /// Create a new RAM machine - pub fn new(config: ConfigArgs) -> Self { + pub fn new(config: ConfigArgs) -> Self { let config = Config::new(K::WORD_SIZE, config); Self { // Memory section @@ -324,7 +324,7 @@ impl_state_machine!(StateMachine); fn main() { // Define the desired machine configuration - let mut machine = StateMachine::::new(DefaultConfig::default_config()); + let mut machine = StateMachine::::new(ConfigArgs::default()); // Show the section map machine.show_sections_maps(); diff --git a/zkmemory/examples/memory-consistency.rs b/zkmemory/examples/memory-consistency.rs index cbaf87f..c220b33 100644 --- a/zkmemory/examples/memory-consistency.rs +++ b/zkmemory/examples/memory-consistency.rs @@ -3,7 +3,7 @@ use rbtree::RBTree; use std::{marker::PhantomData, println}; use zkmemory::{ base::{Base, B256}, - config::{AllocatedSection, Config, ConfigArgs, DefaultConfig}, + config::{Config, ConfigArgs, Section}, constraints::helper::build_and_test_circuit, error::Error, impl_register_machine, impl_stack_machine, impl_state_machine, @@ -54,18 +54,18 @@ where { // Memory memory: RBTree, - memory_allocated: AllocatedSection, + memory_allocated: Section, word_size: K, time_log: u64, // Stack - stack_allocated: AllocatedSection, + stack_allocated: Section, max_stack_depth: u64, stack_depth: u64, stack_ptr: K, // Register - register_allocated: AllocatedSection, + register_allocated: Section, /// Register r0 pub r0: Register, @@ -218,7 +218,7 @@ where V: Base, { /// Create a new RAM machine - pub fn new(config: ConfigArgs) -> Self { + pub fn new(config: ConfigArgs) -> Self { let config = Config::new(K::WORD_SIZE, config); Self { // Memory section @@ -326,7 +326,7 @@ impl_state_machine!(StateMachine); fn main() { // Define the desired machine configuration - let mut machine = StateMachine::::new(DefaultConfig::default_config()); + let mut machine = StateMachine::::new(ConfigArgs::default()); // Show the section map machine.show_sections_maps(); diff --git a/zkmemory/src/config.rs b/zkmemory/src/config.rs index 28288e7..25bfa26 100644 --- a/zkmemory/src/config.rs +++ b/zkmemory/src/config.rs @@ -1,11 +1,14 @@ +extern crate alloc; + use crate::base::Base; use crate::machine::Register; +use alloc::vec::Vec; /// Memory section #[derive(Debug, Clone, Copy)] -pub struct AllocatedSection(T, T); +pub struct Section(T, T); -impl AllocatedSection +impl Section where T: PartialEq + PartialOrd + Copy, { @@ -35,37 +38,87 @@ pub struct Config { /// Buffer size pub buffer_size: T, /// Base address of memory - pub memory: AllocatedSection, + pub memory: Section, /// Stack base address - pub stack: AllocatedSection, + pub stack: Section, /// Register base address - pub register: AllocatedSection, + pub register: Section, } /// Config arguments for RAM machine #[derive(Debug)] -pub struct ConfigArgs { +pub struct ConfigArgs { /// Is head layout pub head_layout: bool, /// Stack depth - pub stack_depth: T, + pub stack_depth: usize, /// Number of registers - pub no_register: T, - /// Buffer size in words - pub buffer_size: T, + pub no_register: usize, + /// Number of buffer elements + pub no_buffer: usize, } -/// Default config -pub struct DefaultConfig; - -impl DefaultConfig { - /// Create a default config - pub fn default_config>() -> ConfigArgs { +impl ConfigArgs { + /// Default configuration + pub fn default() -> Self { ConfigArgs { head_layout: true, - stack_depth: T::from(1024), - no_register: T::from(32), - buffer_size: T::from(32), + stack_depth: 1024, + no_register: 32, + no_buffer: 32, + } + } +} + +/// Memory allocation for RAM machine (Section based implementation) +pub(crate) struct MemoryAllocation { + pub(crate) section: Vec>, + pub(crate) word_size: T, + pub(crate) buffer_size: T, +} + +impl MemoryAllocation +where + T: Base, +{ + /// Create a new memory allocation instance + /// Buffer size is number of elements + pub fn new(word_size: T, no_buffer: usize) -> Self { + MemoryAllocation { + section: Vec::new(), + word_size: word_size, + buffer_size: T::from(no_buffer) * word_size, + } + } + + /// Get the last offset in memory allocation + pub fn last_offset(&self) -> T { + match self.section.last() { + None => T::zero(), + Some(&last) => last.1 + self.buffer_size, + } + } + + /// Add a section to the allocation instance with given size + /// Size is number of elements + pub fn add_fixed_section(&mut self, size: usize) { + let last = self.last_offset(); + self.section + .push(Section(last, last + T::from(size) * self.word_size)); + } + + /// Add a section to the allocation instance with given size + /// Size is number of elements + pub fn add_section(&mut self, end: T) { + let last = self.last_offset(); + self.section.push(Section(last, end)); + } + + /// Set the offset of all sections + pub fn set_offset(&mut self, offset: T) { + for section in &mut self.section { + section.0 = section.0 + offset; + section.1 = section.1 + offset; } } } @@ -75,44 +128,36 @@ where T: Base, { /// Create a new config for given arguments - pub fn new(word_size: T, args: ConfigArgs) -> Self { + pub fn new(word_size: T, args: ConfigArgs) -> Self { + let mut mem_alloc = MemoryAllocation::new(word_size, args.no_buffer); if args.head_layout { - // [stack - buffer - register - buffer - memory] - let stack_lo = T::MIN; - let stack_hi = stack_lo + (args.stack_depth * word_size); - let register_lo = stack_hi + args.buffer_size * word_size; - let register_hi = register_lo + (args.no_register * word_size); - let memory_lo = register_hi + args.buffer_size * word_size; - let memory_hi = T::MAX; + mem_alloc.add_fixed_section(args.stack_depth); + mem_alloc.add_fixed_section(args.no_register); + mem_alloc.add_section(T::MAX); Self { word_size, - stack_depth: args.stack_depth, - buffer_size: args.buffer_size, - stack: AllocatedSection(stack_lo, stack_hi), - register: AllocatedSection(register_lo, register_hi), - memory: AllocatedSection(memory_lo, memory_hi), + stack_depth: T::from(args.stack_depth), + buffer_size: mem_alloc.buffer_size, + stack: mem_alloc.section[0], + register: mem_alloc.section[1], + memory: mem_alloc.section[2], } } else { - // [memory - buffer - stack - buffer - register] - let stack_register_buffer_total_size = - (args.stack_depth + args.no_register + args.buffer_size + args.buffer_size) - * word_size; - let stack_lo = T::MAX - stack_register_buffer_total_size; - let remain = stack_lo % word_size; - let stack_lo = stack_lo - remain; // Align to the nearest previous word-aligned address to ensure sufficient allocation for stack and register sections. - let stack_hi = stack_lo + (args.stack_depth * word_size); - let register_lo = stack_hi + args.buffer_size * word_size; - let register_hi = register_lo + (args.no_register * word_size); - let memory_lo = T::MIN; - let memory_hi = stack_lo - args.buffer_size * word_size; + mem_alloc.add_fixed_section(args.stack_depth); + mem_alloc.add_fixed_section(args.no_register); + let begin_of_fixed_section = T::MAX - mem_alloc.last_offset(); + let begin_of_fixed_section = + begin_of_fixed_section - (begin_of_fixed_section % word_size); + + mem_alloc.set_offset(begin_of_fixed_section); Self { word_size, - stack_depth: args.stack_depth, - buffer_size: args.buffer_size, - stack: AllocatedSection(stack_lo, stack_hi), - register: AllocatedSection(register_lo, register_hi), - memory: AllocatedSection(memory_lo, memory_hi), + stack_depth: T::from(args.stack_depth), + buffer_size: mem_alloc.buffer_size, + stack: mem_alloc.section[0], + register: mem_alloc.section[1], + memory: Section(T::zero(), begin_of_fixed_section - mem_alloc.buffer_size), } } } @@ -130,42 +175,22 @@ where mod tests { use super::ConfigArgs; use crate::base::{Base, B256}; - use crate::config::{Config, DefaultConfig}; - - impl PartialEq for ConfigArgs { - fn eq(&self, other: &Self) -> bool { - self.head_layout == other.head_layout - && self.stack_depth == other.stack_depth - && self.no_register == other.no_register - && self.buffer_size == other.buffer_size - } - } - - #[test] - fn test_default_config() { - let config = ConfigArgs { - head_layout: true, - stack_depth: B256::from(1024), - no_register: B256::from(32), - buffer_size: B256::from(32), - }; - assert_eq!(config, DefaultConfig::default_config()); - } + use crate::config::Config; #[test] fn test_config_sections() { // Test memory section - let config = Config::::new(B256::from(32), DefaultConfig::default_config()); + let config = Config::::new(B256::from(32), ConfigArgs::default()); assert!(config.memory.contain(B256::MAX - B256::from(1))); assert_eq!(config.stack.low(), B256::from(0)); assert_eq!( config.register.low(), - B256::from(config.stack.high() + config.buffer_size * config.word_size) + B256::from(config.stack.high() + config.buffer_size) ); assert_eq!( config.memory.low(), - B256::from(config.register.high() + config.buffer_size * config.word_size) + B256::from(config.register.high() + config.buffer_size) ); assert_eq!(config.memory.high(), B256::MAX); @@ -175,9 +200,9 @@ mod tests { B256::from(32), ConfigArgs { head_layout: false, - stack_depth: B256::from(1024), - no_register, - buffer_size: B256::from(32), + stack_depth: 1024, + no_register: 32, + no_buffer: 32, }, ); @@ -185,11 +210,11 @@ mod tests { assert_eq!(config.memory.low(), B256::from(0)); assert_eq!( config.memory.high(), - B256::from(config.stack.low() - config.buffer_size * config.word_size) + B256::from(config.stack.low() - config.buffer_size) ); assert_eq!( config.stack.high(), - B256::from(config.register.low() - config.buffer_size * config.word_size) + B256::from(config.register.low() - config.buffer_size) ); assert_eq!( config.register.high(), diff --git a/zkmemory/src/machine.rs b/zkmemory/src/machine.rs index 6ca80a5..09999fd 100644 --- a/zkmemory/src/machine.rs +++ b/zkmemory/src/machine.rs @@ -577,7 +577,7 @@ macro_rules! impl_stack_machine { mod tests { use crate::{ base::{Base, B256}, - config::{AllocatedSection, Config, ConfigArgs, DefaultConfig}, + config::{Config, ConfigArgs, Section}, error::Error, machine::{ AbstractContext, AbstractInstruction, AbstractMachine, AbstractMemoryMachine, @@ -631,18 +631,18 @@ mod tests { { // Memory memory: RBTree, - memory_allocated: AllocatedSection, + memory_allocated: Section, word_size: K, time_log: u64, // Stack - stack_allocated: AllocatedSection, + stack_allocated: Section, max_stack_depth: u64, stack_depth: u64, stack_ptr: K, // Register - register_allocated: AllocatedSection, + register_allocated: Section, /// Register r0 pub r0: Register, @@ -804,7 +804,7 @@ mod tests { V: Base, { /// Create a new RAM machine - pub fn new(config: ConfigArgs) -> Self { + pub fn new(config: ConfigArgs) -> Self { let config = Config::new(K::WORD_SIZE, config); Self { // Memory section @@ -923,7 +923,7 @@ mod tests { #[test] fn test_read_write_one_cell() { - let mut sm = StateMachine::::new(DefaultConfig::default_config()); + let mut sm = StateMachine::::new(ConfigArgs::default()); let base = sm.base_address(); let write_chunk = B256::from(1025); let program = vec![ @@ -941,7 +941,7 @@ mod tests { #[test] fn test_read_write_two_cells() { - let mut sm = StateMachine::::new(DefaultConfig::default_config()); + let mut sm = StateMachine::::new(ConfigArgs::default()); let base = sm.base_address(); let write_chunk = [5u8; 32]; let program = vec![ @@ -979,7 +979,7 @@ mod tests { let chunk2 = [190u8; 32]; let add_chunk = [195u8; 32]; - let mut sm = StateMachine::::new(DefaultConfig::default_config()); + let mut sm = StateMachine::::new(ConfigArgs::default()); let base = sm.base_address(); let program = vec![ @@ -1000,7 +1000,7 @@ mod tests { #[test] fn test_stack_machine() { - let mut sm = StateMachine::::new(DefaultConfig::default_config()); + let mut sm = StateMachine::::new(ConfigArgs::default()); assert_eq!(sm.stack_allocated.low(), B256::zero()); let base = sm.base_address(); @@ -1040,7 +1040,7 @@ mod tests { #[test] fn test_stack_machine_part_two() { - let mut sm = StateMachine::::new(DefaultConfig::default_config()); + let mut sm = StateMachine::::new(ConfigArgs::default()); assert_eq!(sm.stack_allocated.low(), B256::zero()); let base = sm.base_address(); @@ -1070,7 +1070,7 @@ mod tests { #[test] #[should_panic] fn test_invalid_instruction() { - let mut sm = StateMachine::::new(DefaultConfig::default_config()); + let mut sm = StateMachine::::new(ConfigArgs::default()); let program = vec![Instruction::Invalid(PhantomData)]; for instruction in program { @@ -1080,7 +1080,7 @@ mod tests { #[test] fn test_register_machine() { - let mut sm = StateMachine::::new(DefaultConfig::default_config()); // by default there are 32 registers allocated + let mut sm = StateMachine::::new(ConfigArgs::default()); // by default there are 32 registers allocated assert_eq!(sm.register_start(), B256::from(33792)); assert_eq!(sm.base_address(), B256::from(35840)); From cb90af698c92f0255bf9b314755c10928c841c89 Mon Sep 17 00:00:00 2001 From: Chiro Hiro Date: Mon, 17 Feb 2025 17:06:26 +0700 Subject: [PATCH 4/5] Fix clippy check and format --- zkmemory/examples/256bits-machine.rs | 2 +- zkmemory/examples/memory-consistency.rs | 2 +- zkmemory/src/config.rs | 6 +++--- zkmemory/src/machine.rs | 14 +++++++------- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/zkmemory/examples/256bits-machine.rs b/zkmemory/examples/256bits-machine.rs index f311c39..8024872 100644 --- a/zkmemory/examples/256bits-machine.rs +++ b/zkmemory/examples/256bits-machine.rs @@ -324,7 +324,7 @@ impl_state_machine!(StateMachine); fn main() { // Define the desired machine configuration - let mut machine = StateMachine::::new(ConfigArgs::default()); + let mut machine = StateMachine::::new(ConfigArgs::default_config()); // Show the section map machine.show_sections_maps(); diff --git a/zkmemory/examples/memory-consistency.rs b/zkmemory/examples/memory-consistency.rs index c220b33..a3836d7 100644 --- a/zkmemory/examples/memory-consistency.rs +++ b/zkmemory/examples/memory-consistency.rs @@ -326,7 +326,7 @@ impl_state_machine!(StateMachine); fn main() { // Define the desired machine configuration - let mut machine = StateMachine::::new(ConfigArgs::default()); + let mut machine = StateMachine::::new(ConfigArgs::default_config()); // Show the section map machine.show_sections_maps(); diff --git a/zkmemory/src/config.rs b/zkmemory/src/config.rs index 25bfa26..d25ddbe 100644 --- a/zkmemory/src/config.rs +++ b/zkmemory/src/config.rs @@ -60,7 +60,7 @@ pub struct ConfigArgs { impl ConfigArgs { /// Default configuration - pub fn default() -> Self { + pub fn default_config() -> Self { ConfigArgs { head_layout: true, stack_depth: 1024, @@ -86,7 +86,7 @@ where pub fn new(word_size: T, no_buffer: usize) -> Self { MemoryAllocation { section: Vec::new(), - word_size: word_size, + word_size, buffer_size: T::from(no_buffer) * word_size, } } @@ -180,7 +180,7 @@ mod tests { #[test] fn test_config_sections() { // Test memory section - let config = Config::::new(B256::from(32), ConfigArgs::default()); + let config = Config::::new(B256::from(32), ConfigArgs::default_config()); assert!(config.memory.contain(B256::MAX - B256::from(1))); assert_eq!(config.stack.low(), B256::from(0)); diff --git a/zkmemory/src/machine.rs b/zkmemory/src/machine.rs index 09999fd..d72f0fd 100644 --- a/zkmemory/src/machine.rs +++ b/zkmemory/src/machine.rs @@ -923,7 +923,7 @@ mod tests { #[test] fn test_read_write_one_cell() { - let mut sm = StateMachine::::new(ConfigArgs::default()); + let mut sm = StateMachine::::new(ConfigArgs::default_config()); let base = sm.base_address(); let write_chunk = B256::from(1025); let program = vec![ @@ -941,7 +941,7 @@ mod tests { #[test] fn test_read_write_two_cells() { - let mut sm = StateMachine::::new(ConfigArgs::default()); + let mut sm = StateMachine::::new(ConfigArgs::default_config()); let base = sm.base_address(); let write_chunk = [5u8; 32]; let program = vec![ @@ -979,7 +979,7 @@ mod tests { let chunk2 = [190u8; 32]; let add_chunk = [195u8; 32]; - let mut sm = StateMachine::::new(ConfigArgs::default()); + let mut sm = StateMachine::::new(ConfigArgs::default_config()); let base = sm.base_address(); let program = vec![ @@ -1000,7 +1000,7 @@ mod tests { #[test] fn test_stack_machine() { - let mut sm = StateMachine::::new(ConfigArgs::default()); + let mut sm = StateMachine::::new(ConfigArgs::default_config()); assert_eq!(sm.stack_allocated.low(), B256::zero()); let base = sm.base_address(); @@ -1040,7 +1040,7 @@ mod tests { #[test] fn test_stack_machine_part_two() { - let mut sm = StateMachine::::new(ConfigArgs::default()); + let mut sm = StateMachine::::new(ConfigArgs::default_config()); assert_eq!(sm.stack_allocated.low(), B256::zero()); let base = sm.base_address(); @@ -1070,7 +1070,7 @@ mod tests { #[test] #[should_panic] fn test_invalid_instruction() { - let mut sm = StateMachine::::new(ConfigArgs::default()); + let mut sm = StateMachine::::new(ConfigArgs::default_config()); let program = vec![Instruction::Invalid(PhantomData)]; for instruction in program { @@ -1080,7 +1080,7 @@ mod tests { #[test] fn test_register_machine() { - let mut sm = StateMachine::::new(ConfigArgs::default()); // by default there are 32 registers allocated + let mut sm = StateMachine::::new(ConfigArgs::default_config()); // by default there are 32 registers allocated assert_eq!(sm.register_start(), B256::from(33792)); assert_eq!(sm.base_address(), B256::from(35840)); From 1b1c83ecfb6b2ef47c4afe1509d27e6b0b0528bd Mon Sep 17 00:00:00 2001 From: Chiro Hiro Date: Mon, 17 Feb 2025 17:10:33 +0700 Subject: [PATCH 5/5] Remove unnecessary config --- zkmemory/Cargo.toml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/zkmemory/Cargo.toml b/zkmemory/Cargo.toml index 4d92154..0a1feae 100644 --- a/zkmemory/Cargo.toml +++ b/zkmemory/Cargo.toml @@ -32,7 +32,3 @@ nova-snark = { workspace = true } bellpepper-core = { workspace = true } poseidon = { path = "../poseidon" } arecibo = { workspace = true } - -[features] -default = ["no_std"] -no_std = [] \ No newline at end of file