Merge branch 'main' of github.com:mir-protocol/plonky2 into ripeMD

evm/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 [dependencies]
 plonky2 = { path = "../plonky2", default-features = false, features = ["rand", "timing"] }
 plonky2_util = { path = "../util" }
-eth-trie-utils = { git = "https://github.com/mir-protocol/eth-trie-utils.git", rev = "3ca443fd18e3f6d209dd96cbad851e05ae058b34" }
+eth-trie-utils = { git = "https://github.com/mir-protocol/eth-trie-utils.git", rev = "c52a04c9f349ac812b886f383a7306b27c8b96dc" }
 maybe_rayon = { path = "../maybe_rayon" }
 anyhow = "1.0.40"
 env_logger = "0.9.0"

evm/src/all_stark.rs

@@ -1,3 +1,5 @@
+use std::iter;
+
 use plonky2::field::extension::Extendable;
 use plonky2::field::types::Field;
 use plonky2::hash::hash_types::RichField;
@@ -5,6 +7,7 @@ use plonky2::hash::hash_types::RichField;
 use crate::config::StarkConfig;
 use crate::cpu::cpu_stark;
 use crate::cpu::cpu_stark::CpuStark;
+use crate::cpu::membus::NUM_GP_CHANNELS;
 use crate::cross_table_lookup::{CrossTableLookup, TableWithColumns};
 use crate::keccak::keccak_stark;
 use crate::keccak::keccak_stark::KeccakStark;
@@ -13,8 +16,8 @@ use crate::keccak_memory::keccak_memory_stark;
 use crate::keccak_memory::keccak_memory_stark::KeccakMemoryStark;
 use crate::logic;
 use crate::logic::LogicStark;
+use crate::memory::memory_stark;
 use crate::memory::memory_stark::MemoryStark;
-use crate::memory::{memory_stark, NUM_CHANNELS};
 use crate::stark::Stark;
 
 #[derive(Clone)]
@@ -129,11 +132,16 @@ fn ctl_logic<F: Field>() -> CrossTableLookup<F> {
 }
 
 fn ctl_memory<F: Field>() -> CrossTableLookup<F> {
-    let cpu_memory_ops = (0..NUM_CHANNELS).map(|channel| {
+    let cpu_memory_code_read = TableWithColumns::new(
+        Table::Cpu,
+        cpu_stark::ctl_data_code_memory(),
+        Some(cpu_stark::ctl_filter_code_memory()),
+    );
+    let cpu_memory_gp_ops = (0..NUM_GP_CHANNELS).map(|channel| {
         TableWithColumns::new(
             Table::Cpu,
-            cpu_stark::ctl_data_memory(channel),
+            cpu_stark::ctl_data_gp_memory(channel),
-            Some(cpu_stark::ctl_filter_memory(channel)),
+            Some(cpu_stark::ctl_filter_gp_memory(channel)),
         )
     });
     let keccak_memory_reads = (0..KECCAK_WIDTH_BYTES).map(|i| {
@@ -150,7 +158,8 @@ fn ctl_memory<F: Field>() -> CrossTableLookup<F> {
             Some(keccak_memory_stark::ctl_filter()),
         )
     });
-    let all_lookers = cpu_memory_ops
+    let all_lookers = iter::once(cpu_memory_code_read)
+        .chain(cpu_memory_gp_ops)
         .chain(keccak_memory_reads)
         .chain(keccak_memory_writes)
         .collect();
@@ -725,6 +734,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore] // Ignoring but not deleting so the test can serve as an API usage example
     fn test_all_stark() -> Result<()> {
         let config = StarkConfig::standard_fast_config();
         let (all_stark, proof) = get_proof(&config)?;
@@ -732,6 +742,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore] // Ignoring but not deleting so the test can serve as an API usage example
     fn test_all_stark_recursive_verifier() -> Result<()> {
         init_logger();
 

evm/src/cpu/columns/mod.rs

@@ -7,6 +7,7 @@ use std::mem::{size_of, transmute};
 use std::ops::{Index, IndexMut};
 
 use crate::cpu::columns::general::CpuGeneralColumnsView;
+use crate::cpu::membus::NUM_GP_CHANNELS;
 use crate::memory;
 use crate::util::{indices_arr, transmute_no_compile_time_size_checks};
 
@@ -35,6 +36,13 @@ pub struct CpuColumnsView<T: Copy> {
     /// Lets us re-use columns in non-cycle rows.
     pub is_cpu_cycle: T,
 
+    /// If CPU cycle: Current context.
+    // TODO: this is currently unconstrained
+    pub context: T,
+
+    /// If CPU cycle: Context for code memory channel.
+    pub code_context: T,
+
     /// If CPU cycle: The program counter for the current instruction.
     pub program_counter: T,
 
@@ -159,7 +167,7 @@ pub struct CpuColumnsView<T: Copy> {
     pub(crate) general: CpuGeneralColumnsView<T>,
 
     pub(crate) clock: T,
-    pub mem_channels: [MemoryChannelView<T>; memory::NUM_CHANNELS],
+    pub mem_channels: [MemoryChannelView<T>; NUM_GP_CHANNELS],
 }
 
 // `u8` is guaranteed to have a `size_of` of 1.

evm/src/cpu/cpu_stark.rs

@@ -1,4 +1,5 @@
 use std::borrow::{Borrow, BorrowMut};
+use std::iter::repeat;
 use std::marker::PhantomData;
 
 use itertools::Itertools;
@@ -10,10 +11,11 @@ use plonky2::hash::hash_types::RichField;
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
 use crate::cpu::columns::{CpuColumnsView, COL_MAP, NUM_CPU_COLUMNS};
 use crate::cpu::{
-    bootstrap_kernel, control_flow, decode, jumps, simple_logic, stack_bounds, syscalls,
+    bootstrap_kernel, control_flow, decode, jumps, membus, simple_logic, stack_bounds, syscalls,
 };
 use crate::cross_table_lookup::Column;
-use crate::memory::NUM_CHANNELS;
+use crate::memory::segments::Segment;
+use crate::memory::{NUM_CHANNELS, VALUE_LIMBS};
 use crate::stark::Stark;
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
@@ -25,14 +27,13 @@ pub fn ctl_data_keccak<F: Field>() -> Vec<Column<F>> {
 }
 
 pub fn ctl_data_keccak_memory<F: Field>() -> Vec<Column<F>> {
-    // When executing KECCAK_GENERAL, the memory channels are used as follows:
+    // When executing KECCAK_GENERAL, the GP memory channels are used as follows:
-    // channel 0: instruction
+    // GP channel 0: stack[-1] = context
-    // channel 1: stack[-1] = context
+    // GP channel 1: stack[-2] = segment
-    // channel 2: stack[-2] = segment
+    // GP channel 2: stack[-3] = virtual
-    // channel 3: stack[-3] = virtual
+    let context = Column::single(COL_MAP.mem_channels[0].value[0]);
-    let context = Column::single(COL_MAP.mem_channels[1].value[0]);
+    let segment = Column::single(COL_MAP.mem_channels[1].value[0]);
-    let segment = Column::single(COL_MAP.mem_channels[2].value[0]);
+    let virt = Column::single(COL_MAP.mem_channels[2].value[0]);
-    let virt = Column::single(COL_MAP.mem_channels[3].value[0]);
 
     let num_channels = F::from_canonical_usize(NUM_CHANNELS);
     let clock = Column::linear_combination([(COL_MAP.clock, num_channels)]);
@@ -60,29 +61,57 @@ pub fn ctl_filter_logic<F: Field>() -> Column<F> {
     Column::sum([COL_MAP.is_and, COL_MAP.is_or, COL_MAP.is_xor])
 }
 
-pub fn ctl_data_memory<F: Field>(channel: usize) -> Vec<Column<F>> {
+pub const MEM_CODE_CHANNEL_IDX: usize = 0;
-    debug_assert!(channel < NUM_CHANNELS);
+pub const MEM_GP_CHANNELS_IDX_START: usize = MEM_CODE_CHANNEL_IDX + 1;
+
+/// Make the time/channel column for memory lookups.
+fn mem_time_and_channel<F: Field>(channel: usize) -> Column<F> {
+    let scalar = F::from_canonical_usize(NUM_CHANNELS);
+    let addend = F::from_canonical_usize(channel);
+    Column::linear_combination_with_constant([(COL_MAP.clock, scalar)], addend)
+}
+
+pub fn ctl_data_code_memory<F: Field>() -> Vec<Column<F>> {
+    let mut cols = vec![
+        Column::constant(F::ONE),                                      // is_read
+        Column::single(COL_MAP.code_context),                          // addr_context
+        Column::constant(F::from_canonical_u64(Segment::Code as u64)), // addr_segment
+        Column::single(COL_MAP.program_counter),                       // addr_virtual
+    ];
+
+    // Low limb of the value matches the opcode bits
+    cols.push(Column::le_bits(COL_MAP.opcode_bits));
+
+    // High limbs of the value are all zero.
+    cols.extend(repeat(Column::constant(F::ZERO)).take(VALUE_LIMBS - 1));
+
+    cols.push(mem_time_and_channel(MEM_CODE_CHANNEL_IDX));
+
+    cols
+}
+
+pub fn ctl_data_gp_memory<F: Field>(channel: usize) -> Vec<Column<F>> {
     let channel_map = COL_MAP.mem_channels[channel];
-    let mut cols: Vec<Column<F>> = Column::singles([
+    let mut cols: Vec<_> = Column::singles([
         channel_map.is_read,
         channel_map.addr_context,
         channel_map.addr_segment,
         channel_map.addr_virtual,
     ])
-    .collect_vec();
+    .collect();
+
     cols.extend(Column::singles(channel_map.value));
 
-    let scalar = F::from_canonical_usize(NUM_CHANNELS);
+    cols.push(mem_time_and_channel(MEM_GP_CHANNELS_IDX_START + channel));
-    let addend = F::from_canonical_usize(channel);
-    cols.push(Column::linear_combination_with_constant(
-        [(COL_MAP.clock, scalar)],
-        addend,
-    ));
 
     cols
 }
 
-pub fn ctl_filter_memory<F: Field>(channel: usize) -> Column<F> {
+pub fn ctl_filter_code_memory<F: Field>() -> Column<F> {
+    Column::single(COL_MAP.is_cpu_cycle)
+}
+
+pub fn ctl_filter_gp_memory<F: Field>(channel: usize) -> Column<F> {
     Column::single(COL_MAP.mem_channels[channel].used)
 }
 
@@ -95,6 +124,7 @@ impl<F: RichField, const D: usize> CpuStark<F, D> {
     pub fn generate(&self, local_values: &mut [F; NUM_CPU_COLUMNS]) {
         let local_values: &mut CpuColumnsView<_> = local_values.borrow_mut();
         decode::generate(local_values);
+        membus::generate(local_values);
         simple_logic::generate(local_values);
         stack_bounds::generate(local_values); // Must come after `decode`.
     }
@@ -117,6 +147,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         control_flow::eval_packed_generic(local_values, next_values, yield_constr);
         decode::eval_packed_generic(local_values, yield_constr);
         jumps::eval_packed(local_values, next_values, yield_constr);
+        membus::eval_packed(local_values, yield_constr);
         simple_logic::eval_packed(local_values, yield_constr);
         stack_bounds::eval_packed(local_values, yield_constr);
         syscalls::eval_packed(local_values, next_values, yield_constr);
@@ -134,6 +165,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         control_flow::eval_ext_circuit(builder, local_values, next_values, yield_constr);
         decode::eval_ext_circuit(builder, local_values, yield_constr);
         jumps::eval_ext_circuit(builder, local_values, next_values, yield_constr);
+        membus::eval_ext_circuit(builder, local_values, yield_constr);
         simple_logic::eval_ext_circuit(builder, local_values, yield_constr);
         stack_bounds::eval_ext_circuit(builder, local_values, yield_constr);
         syscalls::eval_ext_circuit(builder, local_values, next_values, yield_constr);

evm/src/cpu/kernel/aggregator.rs

@@ -47,8 +47,12 @@ pub(crate) fn combined_kernel() -> Kernel {
         include_str!("asm/rlp/encode.asm"),
         include_str!("asm/rlp/decode.asm"),
         include_str!("asm/rlp/read_to_memory.asm"),
-        include_str!("asm/storage/read.asm"),
+        include_str!("asm/mpt/hash.asm"),
-        include_str!("asm/storage/write.asm"),
+        include_str!("asm/mpt/read.asm"),
+        include_str!("asm/mpt/storage_read.asm"),
+        include_str!("asm/mpt/storage_write.asm"),
+        include_str!("asm/mpt/util.asm"),
+        include_str!("asm/mpt/write.asm"),
         include_str!("asm/transactions/router.asm"),
         include_str!("asm/transactions/type_0.asm"),
         include_str!("asm/transactions/type_1.asm"),

evm/src/cpu/kernel/asm/mpt/hash.asm

@@ -0,0 +1,2 @@
+global mpt_hash:
+    // TODO

evm/src/cpu/kernel/asm/mpt/load.asm

@@ -0,0 +1,12 @@
+// Load all partial trie data from prover inputs.
+global mpt_load_all:
+    // First set GLOBAL_METADATA_TRIE_DATA_SIZE = 1.
+    // We don't want it to start at 0, as we use 0 as a null pointer.
+    PUSH 1
+    %mstore(@GLOBAL_METADATA_TRIE_DATA_SIZE)
+
+    TODO
+
+mpt_load_state:
+    PROVER_INPUT(mpt::state)
+    TODO

evm/src/cpu/kernel/asm/mpt/read.asm

@@ -0,0 +1,81 @@
+// Read a value from a MPT.
+//
+// Arguments:
+// - the virtual address of the trie to search in
+// - the key, as a U256
+// - the number of nibbles in the key
+//
+// This function returns a pointer to the leaf, or 0 if the key is not found.
+
+global mpt_read:
+    // stack: node_ptr, key, nibbles, retdest
+    DUP1
+    %mload_trie_data
+    // stack: node_type, node_ptr, key, nibbles, retdest
+    // Increment node_ptr, so it points to the node payload instead of its type.
+    SWAP1 %add_const(1) SWAP1
+    // stack: node_type, node_payload_ptr, key, nibbles, retdest
+
+    DUP1 %eq_const(@MPT_NODE_EMPTY) %jumpi(mpt_read_empty)
+    DUP1 %eq_const(@MPT_NODE_BRANCH) %jumpi(mpt_read_branch)
+    DUP1 %eq_const(@MPT_NODE_EXTENSION) %jumpi(mpt_read_extension)
+    DUP1 %eq_const(@MPT_NODE_LEAF) %jumpi(mpt_read_leaf)
+
+    // There's still the MPT_NODE_DIGEST case, but if we hit a digest node,
+    // it means the prover failed to provide necessary Merkle data, so panic.
+    PANIC
+
+mpt_read_empty:
+    // Return 0 to indicate that the value was not found.
+    %stack (node_type, node_payload_ptr, key, nibbles, retdest)
+        -> (retdest, 0)
+    JUMP
+
+mpt_read_branch:
+    // stack: node_type, node_payload_ptr, key, nibbles, retdest
+    POP
+    // stack: node_payload_ptr, key, nibbles, retdest
+    DUP3 // nibbles
+    ISZERO
+    // stack: nibbles == 0, node_payload_ptr, key, nibbles, retdest
+    %jumpi(mpt_read_branch_end_of_key)
+
+    // stack: node_payload_ptr, key, nibbles, retdest
+    // We have not reached the end of the key, so we descend to one of our children.
+    // Decrement nibbles, then compute current_nibble = (key >> (nibbles * 4)) & 0xF.
+    SWAP2
+    %sub_const(1)
+    // stack: nibbles, key, node_payload_ptr, retdest
+    DUP2 DUP2
+    // stack: nibbles, key, nibbles, key, node_payload_ptr, retdest
+    %mul_const(4)
+    // stack: nibbles * 4, key, nibbles, key, node_payload_ptr, retdest
+    SHR
+    // stack: key >> (nibbles * 4), nibbles, key, node_payload_ptr, retdest
+    %and_const(0xF)
+    // stack: current_nibble, nibbles, key, node_payload_ptr, retdest
+    %stack (current_nibble, nibbles, key, node_payload_ptr, retdest)
+        -> (current_nibble, node_payload_ptr, key, nibbles, retdest)
+    // child_ptr = load(node_payload_ptr + current_nibble)
+    ADD
+    %mload_trie_data
+    // stack: child_ptr, key, nibbles, retdest
+    %jump(mpt_read) // recurse
+
+mpt_read_branch_end_of_key:
+    %stack (node_payload_ptr, key, nibbles, retdest) -> (node_payload_ptr, retdest)
+    // stack: node_payload_ptr, retdest
+    %add_const(16) // skip over the 16 child nodes
+    // stack: leaf_ptr, retdest
+    SWAP1
+    JUMP
+
+mpt_read_extension:
+    // stack: node_type, node_payload_ptr, key, nibbles, retdest
+    POP
+    // stack: node_payload_ptr, key, nibbles, retdest
+
+mpt_read_leaf:
+    // stack: node_type, node_payload_ptr, key, nibbles, retdest
+    POP
+    // stack: node_payload_ptr, key, nibbles, retdest

evm/src/cpu/kernel/asm/mpt/util.asm

@@ -0,0 +1,11 @@
+%macro mload_trie_data
+    // stack: virtual
+    %mload_kernel(@SEGMENT_TRIE_DATA)
+    // stack: value
+%endmacro
+
+%macro mstore_trie_data
+    // stack: virtual, value
+    %mstore_kernel(@SEGMENT_TRIE_DATA)
+    // stack: (empty)
+%endmacro

evm/src/cpu/kernel/asm/mpt/write.asm

@@ -0,0 +1,2 @@
+global mpt_write:
+    // TODO

evm/src/cpu/kernel/assembler.rs

@@ -9,13 +9,13 @@ use crate::cpu::kernel::ast::Item::LocalLabelDeclaration;
 use crate::cpu::kernel::ast::StackReplacement;
 use crate::cpu::kernel::keccak_util::hash_kernel;
 use crate::cpu::kernel::optimizer::optimize_asm;
-use crate::cpu::kernel::prover_input::ProverInputFn;
 use crate::cpu::kernel::stack::stack_manipulation::expand_stack_manipulation;
 use crate::cpu::kernel::utils::u256_to_trimmed_be_bytes;
 use crate::cpu::kernel::{
     ast::{File, Item},
     opcodes::{get_opcode, get_push_opcode},
 };
+use crate::generation::prover_input::ProverInputFn;
 
 /// The number of bytes to push when pushing an offset within the code (i.e. when assembling jumps).
 /// Ideally we would automatically use the minimal number of bytes required, but that would be
@@ -52,6 +52,12 @@ impl Kernel {
     }
 }
 
+#[derive(Eq, PartialEq, Hash, Clone, Debug)]
+struct MacroSignature {
+    name: String,
+    num_params: usize,
+}
+
 struct Macro {
     params: Vec<String>,
     items: Vec<Item>,
@@ -105,17 +111,21 @@ pub(crate) fn assemble(
     Kernel::new(code, global_labels, prover_inputs)
 }
 
-fn find_macros(files: &[File]) -> HashMap<String, Macro> {
+fn find_macros(files: &[File]) -> HashMap<MacroSignature, Macro> {
     let mut macros = HashMap::new();
     for file in files {
         for item in &file.body {
             if let Item::MacroDef(name, params, items) = item {
-                let _macro = Macro {
+                let signature = MacroSignature {
+                    name: name.clone(),
+                    num_params: params.len(),
+                };
+                let macro_ = Macro {
                     params: params.clone(),
                     items: items.clone(),
                 };
-                let old = macros.insert(name.clone(), _macro);
+                let old = macros.insert(signature.clone(), macro_);
-                assert!(old.is_none(), "Duplicate macro: {name}");
+                assert!(old.is_none(), "Duplicate macro signature: {:?}", signature);
             }
         }
     }
@@ -124,7 +134,7 @@ fn find_macros(files: &[File]) -> HashMap<String, Macro> {
 
 fn expand_macros(
     body: Vec<Item>,
-    macros: &HashMap<String, Macro>,
+    macros: &HashMap<MacroSignature, Macro>,
     macro_counter: &mut u32,
 ) -> Vec<Item> {
     let mut expanded = vec![];
@@ -147,30 +157,25 @@ fn expand_macros(
 fn expand_macro_call(
     name: String,
     args: Vec<PushTarget>,
-    macros: &HashMap<String, Macro>,
+    macros: &HashMap<MacroSignature, Macro>,
     macro_counter: &mut u32,
 ) -> Vec<Item> {
-    let _macro = macros
+    let signature = MacroSignature {
-        .get(&name)
-        .unwrap_or_else(|| panic!("No such macro: {}", name));
-
-    assert_eq!(
-        args.len(),
-        _macro.params.len(),
-        "Macro `{}`: expected {} arguments, got {}",
         name,
-        _macro.params.len(),
+        num_params: args.len(),
-        args.len()
+    };
-    );
+    let macro_ = macros
+        .get(&signature)
+        .unwrap_or_else(|| panic!("No such macro: {:?}", signature));
 
     let get_actual_label = |macro_label| format!("@{}.{}", macro_counter, macro_label);
 
     let get_arg = |var| {
-        let param_index = _macro.get_param_index(var);
+        let param_index = macro_.get_param_index(var);
         args[param_index].clone()
     };
 
-    let expanded_item = _macro
+    let expanded_item = macro_
         .items
         .iter()
         .map(|item| match item {
@@ -518,6 +523,18 @@ mod tests {
         parse_and_assemble(&["%macro repeat %endmacro", "%repeat"]);
     }
 
+    #[test]
+    fn overloaded_macros() {
+        let kernel = parse_and_assemble(&[
+            "%macro push(x) PUSH $x %endmacro",
+            "%macro push(x, y) PUSH $x PUSH $y %endmacro",
+            "%push(5)",
+            "%push(6, 7)",
+        ]);
+        let push1 = get_push_opcode(1);
+        assert_eq!(kernel.code, vec![push1, 5, push1, 6, push1, 7]);
+    }
+
     #[test]
     #[should_panic]
     fn macro_with_wrong_vars() {

evm/src/cpu/kernel/ast.rs

@@ -1,6 +1,6 @@
 use ethereum_types::U256;
 
-use crate::cpu::kernel::prover_input::ProverInputFn;
+use crate::generation::prover_input::ProverInputFn;
 
 #[derive(Debug)]
 pub(crate) struct File {

evm/src/cpu/kernel/constants/mod.rs

@@ -4,11 +4,14 @@ use ethereum_types::U256;
 use hex_literal::hex;
 
 use crate::cpu::decode::invalid_opcodes_user;
+use crate::cpu::kernel::constants::trie_type::PartialTrieType;
 use crate::cpu::kernel::context_metadata::ContextMetadata;
 use crate::cpu::kernel::global_metadata::GlobalMetadata;
 use crate::cpu::kernel::txn_fields::NormalizedTxnField;
 use crate::memory::segments::Segment;
 
+pub(crate) mod trie_type;
+
 /// Constants that are accessible to our kernel assembly code.
 pub fn evm_constants() -> HashMap<String, U256> {
     let mut c = HashMap::new();
@@ -30,6 +33,9 @@ pub fn evm_constants() -> HashMap<String, U256> {
     for txn_field in ContextMetadata::all() {
         c.insert(txn_field.var_name().into(), (txn_field as u32).into());
     }
+    for trie_type in PartialTrieType::all() {
+        c.insert(trie_type.var_name().into(), (trie_type as u32).into());
+    }
     c.insert(
         "INVALID_OPCODES_USER".into(),
         U256::from_little_endian(&invalid_opcodes_user()),

evm/src/cpu/kernel/constants/trie_type.rs

@@ -0,0 +1,44 @@
+use eth_trie_utils::partial_trie::PartialTrie;
+
+pub(crate) enum PartialTrieType {
+    Empty = 0,
+    Hash = 1,
+    Branch = 2,
+    Extension = 3,
+    Leaf = 4,
+}
+
+impl PartialTrieType {
+    pub(crate) const COUNT: usize = 5;
+
+    pub(crate) fn of(trie: &PartialTrie) -> Self {
+        match trie {
+            PartialTrie::Empty => Self::Empty,
+            PartialTrie::Hash(_) => Self::Hash,
+            PartialTrie::Branch { .. } => Self::Branch,
+            PartialTrie::Extension { .. } => Self::Extension,
+            PartialTrie::Leaf { .. } => Self::Leaf,
+        }
+    }
+
+    pub(crate) fn all() -> [Self; Self::COUNT] {
+        [
+            Self::Empty,
+            Self::Hash,
+            Self::Branch,
+            Self::Extension,
+            Self::Leaf,
+        ]
+    }
+
+    /// The variable name that gets passed into kernel assembly code.
+    pub(crate) fn var_name(&self) -> &'static str {
+        match self {
+            Self::Empty => "MPT_NODE_EMPTY",
+            Self::Hash => "MPT_NODE_HASH",
+            Self::Branch => "MPT_NODE_BRANCH",
+            Self::Extension => "MPT_NODE_EXTENSION",
+            Self::Leaf => "MPT_NODE_LEAF",
+        }
+    }
+}

evm/src/cpu/kernel/interpreter.rs

@@ -3,14 +3,19 @@ use std::collections::HashMap;
 use anyhow::{anyhow, bail};
 use ethereum_types::{BigEndianHash, U256, U512};
 use keccak_hash::keccak;
+use plonky2::field::goldilocks_field::GoldilocksField;
 
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::assembler::Kernel;
-use crate::cpu::kernel::prover_input::ProverInputFn;
 use crate::cpu::kernel::txn_fields::NormalizedTxnField;
 use crate::generation::memory::{MemoryContextState, MemorySegmentState};
+use crate::generation::prover_input::ProverInputFn;
+use crate::generation::state::GenerationState;
+use crate::generation::GenerationInputs;
 use crate::memory::segments::Segment;
 
+type F = GoldilocksField;
+
 /// Halt interpreter execution whenever a jump to this offset is done.
 const DEFAULT_HALT_OFFSET: usize = 0xdeadbeef;
 
@@ -55,8 +60,8 @@ pub struct Interpreter<'a> {
     offset: usize,
     context: usize,
     pub(crate) memory: InterpreterMemory,
+    generation_state: GenerationState<F>,
     prover_inputs_map: &'a HashMap<usize, ProverInputFn>,
-    prover_inputs: Vec<U256>,
     pub(crate) halt_offsets: Vec<usize>,
     running: bool,
 }
@@ -107,8 +112,8 @@ impl<'a> Interpreter<'a> {
             jumpdests: find_jumpdests(code),
             offset: initial_offset,
             memory: InterpreterMemory::with_code_and_stack(code, initial_stack),
+            generation_state: GenerationState::new(GenerationInputs::default()),
             prover_inputs_map: prover_inputs,
-            prover_inputs: Vec::new(),
             context: 0,
             halt_offsets: vec![DEFAULT_HALT_OFFSET],
             running: true,
@@ -437,9 +442,9 @@ impl<'a> Interpreter<'a> {
             .prover_inputs_map
             .get(&(self.offset - 1))
             .ok_or_else(|| anyhow!("Offset not in prover inputs."))?;
-        let output = prover_input_fn.run(self.stack());
+        let stack = self.stack().to_vec();
+        let output = self.generation_state.prover_input(&stack, prover_input_fn);
         self.push(output);
-        self.prover_inputs.push(output);
         Ok(())
     }
 

evm/src/cpu/kernel/mod.rs

@@ -1,7 +1,7 @@
 pub mod aggregator;
 pub mod assembler;
 mod ast;
-mod constants;
+pub(crate) mod constants;
 pub(crate) mod context_metadata;
 mod cost_estimator;
 pub(crate) mod global_metadata;
@@ -9,7 +9,6 @@ pub(crate) mod keccak_util;
 mod opcodes;
 mod optimizer;
 mod parser;
-pub mod prover_input;
 pub mod stack;
 mod txn_fields;
 mod utils;

evm/src/cpu/membus.rs

@@ -0,0 +1,70 @@
+use plonky2::field::extension::Extendable;
+use plonky2::field::packed::PackedField;
+use plonky2::field::types::PrimeField64;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::ext_target::ExtensionTarget;
+
+use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+use crate::cpu::columns::CpuColumnsView;
+
+/// General-purpose memory channels; they can read and write to all contexts/segments/addresses.
+pub const NUM_GP_CHANNELS: usize = 4;
+
+pub mod channel_indices {
+    use std::ops::Range;
+
+    pub const CODE: usize = 0;
+    pub const GP: Range<usize> = CODE + 1..(CODE + 1) + super::NUM_GP_CHANNELS;
+}
+
+/// Total memory channels used by the CPU table. This includes all the `GP_MEM_CHANNELS` as well as
+/// all special-purpose memory channels.
+///
+/// Currently, there is one special-purpose memory channel, which reads the opcode from memory. Its
+/// limitations are:
+///  - it is enabled by `is_cpu_cycle`,
+///  - it always reads and cannot write,
+///  - the context is derived from the current context and the `is_kernel_mode` flag,
+///  - the segment is hard-wired to the code segment,
+///  - the address is `program_counter`,
+///  - the value must fit in one byte (in the least-significant position) and its eight bits are
+///    found in `opcode_bits`.
+/// These limitations save us numerous columns in the CPU table.
+pub const NUM_CHANNELS: usize = channel_indices::GP.end;
+
+/// Calculates `lv.stack_len_bounds_aux`. Note that this must be run after decode.
+pub fn generate<F: PrimeField64>(lv: &mut CpuColumnsView<F>) {
+    let cycle_filter = lv.is_cpu_cycle;
+    if cycle_filter == F::ZERO {
+        return;
+    }
+
+    assert!(lv.is_kernel_mode.to_canonical_u64() <= 1);
+
+    // Set `lv.code_context` to 0 if in kernel mode and to `lv.context` if in user mode.
+    lv.code_context = (F::ONE - lv.is_kernel_mode) * lv.context;
+}
+
+pub fn eval_packed<P: PackedField>(
+    lv: &CpuColumnsView<P>,
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    // Validate `lv.code_context`. It should be 0 if in kernel mode and `lv.context` if in user
+    // mode.
+    yield_constr.constraint(
+        lv.is_cpu_cycle * (lv.code_context - (P::ONES - lv.is_kernel_mode) * lv.context),
+    );
+}
+
+pub fn eval_ext_circuit<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut plonky2::plonk::circuit_builder::CircuitBuilder<F, D>,
+    lv: &CpuColumnsView<ExtensionTarget<D>>,
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    // Validate `lv.code_context`. It should be 0 if in kernel mode and `lv.context` if in user
+    // mode.
+    let diff = builder.sub_extension(lv.context, lv.code_context);
+    let constr = builder.mul_sub_extension(lv.is_kernel_mode, lv.context, diff);
+    let filtered_constr = builder.mul_extension(lv.is_cpu_cycle, constr);
+    yield_constr.constraint(builder, filtered_constr);
+}

evm/src/cpu/mod.rs

@@ -5,6 +5,7 @@ pub mod cpu_stark;
 pub(crate) mod decode;
 mod jumps;
 pub mod kernel;
+pub(crate) mod membus;
 mod simple_logic;
 mod stack_bounds;
 mod syscalls;

evm/src/generation/mod.rs

@@ -20,9 +20,11 @@ use crate::proof::{BlockMetadata, PublicValues, TrieRoots};
 use crate::util::trace_rows_to_poly_values;
 
 pub(crate) mod memory;
+pub(crate) mod mpt;
+pub(crate) mod prover_input;
 pub(crate) mod state;
 
-#[derive(Clone, Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize, Default)]
 /// Inputs needed for trace generation.
 pub struct GenerationInputs {
     pub signed_txns: Vec<Vec<u8>>,
@@ -55,7 +57,7 @@ pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     inputs: GenerationInputs,
     config: &StarkConfig,
 ) -> ([Vec<PolynomialValues<F>>; NUM_TABLES], PublicValues) {
-    let mut state = GenerationState::<F>::default();
+    let mut state = GenerationState::<F>::new(inputs.clone());
 
     generate_bootstrap_kernel::<F>(&mut state);
 

evm/src/generation/mpt.rs

@@ -0,0 +1,61 @@
+use eth_trie_utils::partial_trie::PartialTrie;
+use ethereum_types::U256;
+
+use crate::cpu::kernel::constants::trie_type::PartialTrieType;
+use crate::generation::GenerationInputs;
+
+pub(crate) fn all_mpt_prover_inputs(inputs: &GenerationInputs) -> Vec<U256> {
+    let mut prover_inputs = vec![];
+
+    mpt_prover_inputs(&inputs.state_trie, &mut prover_inputs, &|_rlp| vec![]); // TODO
+
+    mpt_prover_inputs(
+        &inputs.transactions_trie,
+        &mut prover_inputs,
+        &|_rlp| vec![],
+    ); // TODO
+
+    mpt_prover_inputs(&inputs.receipts_trie, &mut prover_inputs, &|_rlp| vec![]); // TODO
+
+    for (_addr, storage_trie) in &inputs.storage_tries {
+        mpt_prover_inputs(storage_trie, &mut prover_inputs, &|leaf_be| {
+            vec![U256::from_big_endian(leaf_be)]
+        });
+    }
+
+    prover_inputs
+}
+
+/// Given a trie, generate the prover input data for that trie. In essence, this serializes a trie
+/// into a `U256` array, in a simple format which the kernel understands. For example, a leaf node
+/// is serialized as `(TYPE_LEAF, key, value)`, where key is a `(nibbles, depth)` pair and `value`
+/// is a variable-length structure which depends on which trie we're dealing with.
+pub(crate) fn mpt_prover_inputs<F>(
+    trie: &PartialTrie,
+    prover_inputs: &mut Vec<U256>,
+    parse_leaf: &F,
+) where
+    F: Fn(&[u8]) -> Vec<U256>,
+{
+    prover_inputs.push((PartialTrieType::of(trie) as u32).into());
+    match trie {
+        PartialTrie::Empty => {}
+        PartialTrie::Hash(h) => prover_inputs.push(*h),
+        PartialTrie::Branch { children, value } => {
+            for child in children {
+                mpt_prover_inputs(child, prover_inputs, parse_leaf);
+            }
+            prover_inputs.extend(parse_leaf(value));
+        }
+        PartialTrie::Extension { nibbles, child } => {
+            prover_inputs.push(nibbles.count.into());
+            prover_inputs.push(nibbles.packed);
+            mpt_prover_inputs(child, prover_inputs, parse_leaf);
+        }
+        PartialTrie::Leaf { nibbles, value } => {
+            prover_inputs.push(nibbles.count.into());
+            prover_inputs.push(nibbles.packed);
+            prover_inputs.extend(parse_leaf(value));
+        }
+    }
+}

evm/src/generation/prover_input.rs

@@ -1,11 +1,13 @@
 use std::str::FromStr;
 
 use ethereum_types::U256;
+use plonky2::field::types::Field;
 
-use crate::cpu::kernel::prover_input::Field::{
+use crate::generation::prover_input::EvmField::{
     Bn254Base, Bn254Scalar, Secp256k1Base, Secp256k1Scalar,
 };
-use crate::cpu::kernel::prover_input::FieldOp::{Inverse, Sqrt};
+use crate::generation::prover_input::FieldOp::{Inverse, Sqrt};
+use crate::generation::state::GenerationState;
 
 /// Prover input function represented as a scoped function name.
 /// Example: `PROVER_INPUT(ff::bn254_base::inverse)` is represented as `ProverInputFn([ff, bn254_base, inverse])`.
@@ -18,32 +20,39 @@ impl From<Vec<String>> for ProverInputFn {
     }
 }
 
-impl ProverInputFn {
+impl<F: Field> GenerationState<F> {
-    /// Run the function on the stack.
+    #[allow(unused)] // TODO: Should be used soon.
-    pub fn run(&self, stack: &[U256]) -> U256 {
+    pub(crate) fn prover_input(&mut self, stack: &[U256], input_fn: &ProverInputFn) -> U256 {
-        match self.0[0].as_str() {
+        match input_fn.0[0].as_str() {
-            "ff" => self.run_ff(stack),
+            "ff" => self.run_ff(stack, input_fn),
-            "mpt" => todo!(),
+            "mpt" => self.run_mpt(input_fn),
             _ => panic!("Unrecognized prover input function."),
         }
     }
 
-    // Finite field operations.
+    /// Finite field operations.
-    fn run_ff(&self, stack: &[U256]) -> U256 {
+    fn run_ff(&self, stack: &[U256], input_fn: &ProverInputFn) -> U256 {
-        let field = Field::from_str(self.0[1].as_str()).unwrap();
+        let field = EvmField::from_str(input_fn.0[1].as_str()).unwrap();
-        let op = FieldOp::from_str(self.0[2].as_str()).unwrap();
+        let op = FieldOp::from_str(input_fn.0[2].as_str()).unwrap();
         let x = *stack.last().expect("Empty stack");
         field.op(op, x)
     }
 
-    // MPT operations.
+    /// MPT data.
-    #[allow(dead_code)]
+    fn run_mpt(&mut self, input_fn: &ProverInputFn) -> U256 {
-    fn run_mpt(&self, _stack: Vec<U256>) -> U256 {
+        let operation = input_fn.0[1].as_str();
-        todo!()
+        match operation {
+            "trie_data" => self
+                .mpt_prover_inputs
+                .pop()
+                .unwrap_or_else(|| panic!("Out of MPT data")),
+            "num_storage_tries" => self.inputs.storage_tries.len().into(),
+            _ => panic!("Unrecognized MPT operation."),
+        }
     }
 }
 
-enum Field {
+enum EvmField {
     Bn254Base,
     Bn254Scalar,
     Secp256k1Base,
@@ -55,7 +64,7 @@ enum FieldOp {
     Sqrt,
 }
 
-impl FromStr for Field {
+impl FromStr for EvmField {
     type Err = ();
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
@@ -81,19 +90,19 @@ impl FromStr for FieldOp {
     }
 }
 
-impl Field {
+impl EvmField {
     fn order(&self) -> U256 {
         match self {
-            Field::Bn254Base => {
+            EvmField::Bn254Base => {
                 U256::from_str("0x30644e72e131a029b85045b68181585d97816a916871ca8d3c208c16d87cfd47")
                     .unwrap()
             }
-            Field::Bn254Scalar => todo!(),
+            EvmField::Bn254Scalar => todo!(),
-            Field::Secp256k1Base => {
+            EvmField::Secp256k1Base => {
                 U256::from_str("0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f")
                     .unwrap()
             }
-            Field::Secp256k1Scalar => {
+            EvmField::Secp256k1Scalar => {
                 U256::from_str("0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141")
                     .unwrap()
             }

evm/src/generation/state.rs

@@ -6,6 +6,8 @@ use tiny_keccak::keccakf;
 
 use crate::cpu::columns::{CpuColumnsView, NUM_CPU_COLUMNS};
 use crate::generation::memory::MemoryState;
+use crate::generation::mpt::all_mpt_prover_inputs;
+use crate::generation::GenerationInputs;
 use crate::keccak_memory::keccak_memory_stark::KeccakMemoryOp;
 use crate::memory::memory_stark::MemoryOp;
 use crate::memory::segments::Segment;
@@ -15,6 +17,7 @@ use crate::{keccak, logic};
 
 #[derive(Debug)]
 pub(crate) struct GenerationState<F: Field> {
+    pub(crate) inputs: GenerationInputs,
     pub(crate) cpu_rows: Vec<[F; NUM_CPU_COLUMNS]>,
     pub(crate) current_cpu_row: CpuColumnsView<F>,
 
@@ -24,9 +27,30 @@ pub(crate) struct GenerationState<F: Field> {
     pub(crate) keccak_inputs: Vec<[u64; keccak::keccak_stark::NUM_INPUTS]>,
     pub(crate) keccak_memory_inputs: Vec<KeccakMemoryOp>,
     pub(crate) logic_ops: Vec<logic::Operation>,
+
+    /// Prover inputs containing MPT data, in reverse order so that the next input can be obtained
+    /// via `pop()`.
+    pub(crate) mpt_prover_inputs: Vec<U256>,
 }
 
 impl<F: Field> GenerationState<F> {
+    pub(crate) fn new(inputs: GenerationInputs) -> Self {
+        let mut mpt_prover_inputs = all_mpt_prover_inputs(&inputs);
+        mpt_prover_inputs.reverse();
+
+        Self {
+            inputs,
+            cpu_rows: vec![],
+            current_cpu_row: [F::ZERO; NUM_CPU_COLUMNS].into(),
+            current_context: 0,
+            memory: MemoryState::default(),
+            keccak_inputs: vec![],
+            keccak_memory_inputs: vec![],
+            logic_ops: vec![],
+            mpt_prover_inputs,
+        }
+    }
+
     /// Compute logical AND, and record the operation to be added in the logic table later.
     #[allow(unused)] // TODO: Should be used soon.
     pub(crate) fn and(&mut self, input0: U256, input1: U256) -> U256 {
@@ -217,19 +241,3 @@ impl<F: Field> GenerationState<F> {
         self.cpu_rows.push(swapped_row.into());
     }
 }
-
-// `GenerationState` can't `derive(Default)` because `Default` is only implemented for arrays up to
-// length 32 :-\.
-impl<F: Field> Default for GenerationState<F> {
-    fn default() -> Self {
-        Self {
-            cpu_rows: vec![],
-            current_cpu_row: [F::ZERO; NUM_CPU_COLUMNS].into(),
-            current_context: 0,
-            memory: MemoryState::default(),
-            keccak_inputs: vec![],
-            keccak_memory_inputs: vec![],
-            logic_ops: vec![],
-        }
-    }
-}

evm/src/memory/mod.rs

@@ -3,5 +3,5 @@ pub mod memory_stark;
 pub mod segments;
 
 // TODO: Move to CPU module, now that channels have been removed from the memory table.
-pub(crate) const NUM_CHANNELS: usize = 4;
+pub(crate) const NUM_CHANNELS: usize = crate::cpu::membus::NUM_CHANNELS;
 pub(crate) const VALUE_LIMBS: usize = 8;