Standard VM customization

In this section, we will guide you through the key components of a memory-verifiable VM built with zkMemory and how to customize them. By the end of this tour, you will understand how to tailor your VM configuration, execute a program, and generate proofs.

1. Memory Layout in zkMemory

The memory layout in zkMemory is a critical aspect of the system, as it dictates how different components of the VM (such as memory, registers, and stack) are organized in memory. The layout is configurable, allowing for both head and tail configurations, and consists of several key sections:

• Stack Section: This is where the stack data is stored. The stack uses a special area of memory that is dynamically accessed and modified.
• Buffer Section: A temporary buffer section used to prevent out-of-bounds memory access.
• Register Section: Reserved for storing the register values.
• Memory Section: The main memory used for data storage.

You can choose between a head layout (stack at the top) or tail layout (stack at the bottom) when configuring the VM.

Illustration of Memory Layouts:

Head Layout:

┌──────────────────┐ ┌──────┐ ┌──────────────────┐ ┌──────┐ ┌──────────────────┐
│  Stack Section   │ │Buffer│ │ Register Section │ │Buffer│ │  Memory Section  │
└──────────────────┘ └──────┘ └──────────────────┘ └──────┘ └──────────────────┘

Tail Layout:

┌──────────────────┐ ┌──────┐ ┌──────────────────┐ ┌──────┐ ┌──────────────────┐
│  Memory Section  │ │Buffer│ │  Stack Section   │ │Buffer│ │ Register Section │
└──────────────────┘ └──────┘ └──────────────────┘ └──────┘ └──────────────────┘

2. Customizing Memory Layout

To customize the layout, we need to define the layout type, word size, stack depth, and buffer size. These are configurable via the ConfigArgs struct, which is used to craft the Config object that holds the final configuration, including important information such as address range for each section.

Here's the structure of the Config and ConfigArgs struct which you can use to customize your VM layout:

/// Config for RAM machine
#[derive(Debug, Clone, Copy)]
pub struct Config<T, const S: usize> {
    pub word_size: T,
    pub stack_depth: T,
    pub buffer_size: T,
    pub memory: AllocatedSection<T>,
    pub stack: AllocatedSection<T>,
    pub register: AllocatedSection<T>,
}

#[derive(Debug)]
pub struct ConfigArgs<T> {
    /// Is head layout
    pub head_layout: bool,
    /// Stack depth
    pub stack_depth: T,
    /// Number of registers
    pub no_register: T,
    /// Buffer size
    pub buffer_size: T,
}

/// Default config
pub struct DefaultConfig;

impl DefaultConfig {
    pub fn default_config<const S: usize, T: Base<S>>() -> ConfigArgs<T> {
        ConfigArgs {
            head_layout: true,
            stack_depth: T::from(1024),
            no_register: T::from(32),
            buffer_size: T::from(32),
        }
    }
}
impl<T, const S: usize> Config<T, S>
where
    T: Base<S>,
{
    /// Create a new config for given config arguments
    pub fn new(word_size: T, args: ConfigArgs<T>) -> Self;
}

In this configuration:

• The word_size determines the size of a memory cell (e.g., U256, u64, u32).
• The stack_depth controls how deep the stack can grow.
• The buffer_size prevents out-of-bounds memory access by reserving space.

To create a configuration with specific word sizes and other parameters, you can use:

let config = Config::new(U256::from(256), ConfigArgs {
    stack_depth: 1024,
    buffer_size: 32,
    head_layout: true,
    no_register: 32,
});

3. Creating a standard VM

Once the configuration is set, you can declare the VM and the memory sections using the Config struct. This is done through the StandardStateMachine struct which implements the AbstractMachine trait:

type CustomStateMachine = StandardStateMachine<B256, B256, 32, 32>;
let mut machine = CustomStateMachine::new(DefaultConfig::default_config());

In the standard state machine version implemented in StandardStateMachine and StandardInstruction, we use a Red-black tree to store memory data. Therefore the data accesses are represented as tuples in the format (key, value). The first two params of StandardStateMachine represents the numeric type for the keys and values and the last two constant number params explicitly specify the number of bytes in the former two types respectively. When using the default config (DefaultConfig::default_config()), the word size is 32 bits by default.

The machine object now has access to the configured memory, stack, and register sections. You can print the section map to verify the layout:

let sections = machine.get_sections_maps();
for (i, (start, end)) in sections.iter().enumerate() {
    let section_name = match i {
        0 => "Memory",
        1 => "Register",
        2 => "Stack",
        _ => "Unknown",
    };
    println!("{}: ({}, {})", section_name, start, end);
}

4. Specifying Program Execution (Trace)

To execute a program, you need to list the instructions in a sequence. Instructions are defined in the VM’s architecture, and can include operations like WRITE, READ, PUSH, POP, etc.

Here’s an example of how to define and execute a program:

// declare alias for standard instruction type, and it should be compatible with the type of the standard state machine being used
type Instruction = StandardInstruction<CustomStateMachine, B256, B256, 32, 32>;

// specify the list of instruction executions sequentially
let program = vec![
    Instruction::Write(base + B256::from(16), B256::from(randomize.gen_range(u64::MAX / 2..u64::MAX))),
    Instruction::Write(base + B256::from(48), B256::from(randomize.gen_range(u64::MAX / 2..u64::MAX))),
    Instruction::Read(base + B256::from(16)),
    Instruction::Push(B256::from(777)),
    Instruction::Swap(machine.r0),
    Instruction::Mov(machine.r1, machine.r0),
    Instruction::Write(base + B256::from(10000), B256::from(randomize.gen_range(u64::MAX / 2..u64::MAX))),
];

The VM will execute these instructions sequentially, and you can debug the execution by printing the trace of each instruction:

for instruction in program {
    println!("Instruction: {:?}", instruction);
    machine.exec(&instruction);
}

After execution, you can print the trace record:

for x in machine.trace().into_iter() {
    println!("{:?}", x);
}

5. Proofs generation and verification

Finally, once the program is executed, you can generate proofs to verify the correctness of the VM’s behavior. zkMemory integrates deeply with the Halo2 framework for proof generation and verification.

let log2_number_of_ir_value_rows = 10;
let sorted_trace = sort_trace::<B256, B256, 32, 32>(trace_record.clone());

let circuit = MemoryConsistencyCircuit::<Fp> {
    input: trace_record.clone(),
    shuffle: sorted_trace.clone(),
    marker: PhantomData,
};

let prover = MockProver::run(log2_number_of_ir_value_rows, &circuit, vec![])
    .expect("Cannot run the circuit");
assert_eq!(prover.verify(), Ok(()));

Note that the parameter log2_number_of_ir_value_rows function represents the ceiling of the log2 of the number of rows in the compiled value table in Halo2. Therefore, you should increase this parameter log2_number_of_ir_value_rows if it is not large enough and incurs errors.

If no errors occur, the trace is valid and you can proceed with the proof verification. This ensures that the VM’s memory operations are consistent and verifiable through zero-knowledge proofs.

Optionally, you can also replace this part with the method build_and_test_circuit() for brevity.

// function signature of the proof generation function
pub fn build_and_test_circuit(trace: Vec<TraceRecord<B256, B256, 32, 32>>, k: u32); // k is log2_number_of_ir_value_rows

// calling the proof generation function
build_and_test_circuit(trace_record, 10);

---

By following these steps, you can create a customized, memory-verifiable VM in zkMemory, define a program, and generate proofs for its correctness.