Implement Polygon Hermez type2 SMT

Cargo.toml

@@ -1,5 +1,11 @@
 [workspace]
-members = ["mpt_trie", "proof_gen", "trace_decoder", "evm_arithmetization"]
+members = [
+    "evm_arithmetization",
+    "mpt_trie",
+    "proof_gen",
+    "smt_trie",
+    "trace_decoder"
+]
 resolver = "2"
 
 [workspace.dependencies]

evm_arithmetization/Cargo.toml

@@ -42,6 +42,7 @@ serde_json = { workspace = true }
 
 # Local dependencies
 mpt_trie = { version = "0.1.0", path = "../mpt_trie" }
+smt_trie = { version = "0.1.0", path = "../smt_trie" }
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 jemallocator = "0.5.0"

evm_arithmetization/src/all_stark.rs

@@ -23,6 +23,8 @@ use crate::logic;
 use crate::logic::LogicStark;
 use crate::memory::memory_stark;
 use crate::memory::memory_stark::MemoryStark;
+use crate::poseidon::poseidon_stark;
+use crate::poseidon::poseidon_stark::PoseidonStark;
 
 /// Structure containing all STARKs and the cross-table lookups.
 #[derive(Clone)]
@@ -34,6 +36,7 @@ pub struct AllStark<F: RichField + Extendable<D>, const D: usize> {
     pub(crate) keccak_sponge_stark: KeccakSpongeStark<F, D>,
     pub(crate) logic_stark: LogicStark<F, D>,
     pub(crate) memory_stark: MemoryStark<F, D>,
+    pub(crate) poseidon_stark: PoseidonStark<F, D>,
     pub(crate) cross_table_lookups: Vec<CrossTableLookup<F>>,
 }
 
@@ -49,6 +52,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Default for AllStark<F, D> {
             keccak_sponge_stark: KeccakSpongeStark::default(),
             logic_stark: LogicStark::default(),
             memory_stark: MemoryStark::default(),
+            poseidon_stark: PoseidonStark::default(),
             cross_table_lookups: all_cross_table_lookups(),
         }
     }
@@ -64,6 +68,7 @@ impl<F: RichField + Extendable<D>, const D: usize> AllStark<F, D> {
             self.keccak_sponge_stark.num_lookup_helper_columns(config),
             self.logic_stark.num_lookup_helper_columns(config),
             self.memory_stark.num_lookup_helper_columns(config),
+            0,
         ]
     }
 }
@@ -80,6 +85,7 @@ pub enum Table {
     KeccakSponge = 4,
     Logic = 5,
     Memory = 6,
+    Poseidon = 7,
 }
 
 impl Deref for Table {
@@ -88,12 +94,12 @@ impl Deref for Table {
     fn deref(&self) -> &Self::Target {
         // Hacky way to implement `Deref` for `Table` so that we don't have to
         // call `Table::Foo as usize`, but perhaps too ugly to be worth it.
-        [&0, &1, &2, &3, &4, &5, &6][*self as TableIdx]
+        [&0, &1, &2, &3, &4, &5, &6, &7][*self as TableIdx]
     }
 }
 
 /// Number of STARK tables.
-pub(crate) const NUM_TABLES: usize = Table::Memory as usize + 1;
+pub(crate) const NUM_TABLES: usize = Table::Poseidon as usize + 1;
 
 impl Table {
     /// Returns all STARK table indices.
@@ -106,6 +112,7 @@ impl Table {
             Self::KeccakSponge,
             Self::Logic,
             Self::Memory,
+            Self::Poseidon,
         ]
     }
 }
@@ -120,6 +127,7 @@ pub(crate) fn all_cross_table_lookups<F: Field>() -> Vec<CrossTableLookup<F>> {
         ctl_keccak_outputs(),
         ctl_logic(),
         ctl_memory(),
+        ctl_poseidon(),
     ]
 }
 
@@ -306,3 +314,10 @@ fn ctl_memory<F: Field>() -> CrossTableLookup<F> {
     );
     CrossTableLookup::new(all_lookers, memory_looked)
 }
+
+fn ctl_poseidon<F: Field>() -> CrossTableLookup<F> {
+    CrossTableLookup::new(
+        vec![cpu_stark::ctl_poseidon()],
+        poseidon_stark::ctl_looked(),
+    )
+}

evm_arithmetization/src/cpu/columns/ops.rs

@@ -24,6 +24,7 @@ pub(crate) struct OpsColumnsView<T: Copy> {
     pub shift: T,
     /// Combines JUMPDEST and KECCAK_GENERAL flags.
     pub jumpdest_keccak_general: T,
+    pub poseidon: T,
     /// Combines JUMP and JUMPI flags.
     pub jumps: T,
     /// Combines PUSH and PROVER_INPUT flags.

evm_arithmetization/src/cpu/contextops.rs

@@ -23,6 +23,7 @@ const KEEPS_CONTEXT: OpsColumnsView<bool> = OpsColumnsView {
     not_pop: true,
     shift: true,
     jumpdest_keccak_general: true,
+    poseidon: true,
     push_prover_input: true,
     jumps: true,
     pc_push0: true,

evm_arithmetization/src/cpu/control_flow.rs

@@ -8,7 +8,7 @@ use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsume
 use crate::cpu::columns::{CpuColumnsView, COL_MAP};
 use crate::cpu::kernel::aggregator::KERNEL;
 
-const NATIVE_INSTRUCTIONS: [usize; 12] = [
+const NATIVE_INSTRUCTIONS: [usize; 13] = [
     COL_MAP.op.binary_op,
     COL_MAP.op.ternary_op,
     COL_MAP.op.fp254_op,
@@ -17,6 +17,7 @@ const NATIVE_INSTRUCTIONS: [usize; 12] = [
     COL_MAP.op.not_pop,
     COL_MAP.op.shift,
     COL_MAP.op.jumpdest_keccak_general,
+    COL_MAP.op.poseidon,
     // Not PROVER_INPUT: it is dealt with manually below.
     // not JUMPS (possible need to jump)
     COL_MAP.op.pc_push0,

evm_arithmetization/src/cpu/cpu_stark.rs

@@ -429,6 +429,31 @@ pub(crate) fn ctl_filter_set_context<F: Field>() -> Filter<F> {
     )
 }
 
+/// Returns the `TableWithColumns` for the CPU rows calling POSEIDON.
+pub(crate) fn ctl_poseidon<F: Field>() -> TableWithColumns<F> {
+    let mut columns = Vec::new();
+    for channel in 0..3 {
+        for i in 0..VALUE_LIMBS / 2 {
+            columns.push(Column::linear_combination([
+                (COL_MAP.mem_channels[channel].value[2 * i], F::ONE),
+                (
+                    COL_MAP.mem_channels[channel].value[2 * i + 1],
+                    F::from_canonical_u64(1 << 32),
+                ),
+            ]));
+        }
+    }
+    columns.extend(Column::singles_next_row(COL_MAP.mem_channels[0].value));
+    // res.extend(Column::singles(COL_MAP.mem_channels[1].value));
+    // res.extend(Column::singles(COL_MAP.mem_channels[2].value));
+    // res
+    TableWithColumns::new(
+        *Table::Cpu,
+        columns,
+        Some(Filter::new_simple(Column::single(COL_MAP.op.poseidon))),
+    )
+}
+
 /// Disable the specified memory channels.
 /// Since channel 0 contains the top of the stack and is handled specially,
 /// channels to disable are 1, 2 or both. All cases can be expressed as a vec.

evm_arithmetization/src/cpu/decode.rs

@@ -25,13 +25,14 @@ use crate::cpu::columns::{CpuColumnsView, COL_MAP};
 /// Note: invalid opcodes are not represented here. _Any_ opcode is permitted to
 /// decode to `is_invalid`. The kernel then verifies that the opcode was
 /// _actually_ invalid.
-const OPCODES: [(u8, usize, bool, usize); 5] = [
+const OPCODES: [(u8, usize, bool, usize); 6] = [
     // (start index of block, number of top bits to check (log2), kernel-only, flag column)
     // ADD, MUL, SUB, DIV, MOD, LT, GT and BYTE flags are handled partly manually here, and partly
     // through the Arithmetic table CTL. ADDMOD, MULMOD and SUBMOD flags are handled partly
     // manually here, and partly through the Arithmetic table CTL. FP254 operation flags are
     // handled partly manually here, and partly through the Arithmetic table CTL.
     (0x14, 1, false, COL_MAP.op.eq_iszero),
+    (0x22, 0, true, COL_MAP.op.poseidon),
     // AND, OR and XOR flags are handled partly manually here, and partly through the Logic table
     // CTL. NOT and POP are handled manually here.
     // SHL and SHR flags are handled partly manually here, and partly through the Logic table CTL.

evm_arithmetization/src/cpu/gas.rs

@@ -27,8 +27,9 @@ const SIMPLE_OPCODES: OpsColumnsView<Option<u32>> = OpsColumnsView {
     not_pop: None, // This is handled manually below
     shift: G_VERYLOW,
     jumpdest_keccak_general: None, // This is handled manually below.
-    push_prover_input: None,       // This is handled manually below.
-    jumps: None,                   // Combined flag handled separately.
+    poseidon: KERNEL_ONLY_INSTR,
+    push_prover_input: None, // This is handled manually below.
+    jumps: None,             // Combined flag handled separately.
     pc_push0: G_BASE,
     dup_swap: G_VERYLOW,
     context_op: KERNEL_ONLY_INSTR,

evm_arithmetization/src/cpu/kernel/aggregator.rs

@@ -126,6 +126,12 @@ pub(crate) fn combined_kernel() -> Kernel {
         include_str!("asm/mpt/storage/storage_read.asm"),
         include_str!("asm/mpt/storage/storage_write.asm"),
         include_str!("asm/mpt/util.asm"),
+        include_str!("asm/smt/delete.asm"),
+        include_str!("asm/smt/hash.asm"),
+        include_str!("asm/smt/insert.asm"),
+        include_str!("asm/smt/keys.asm"),
+        include_str!("asm/smt/read.asm"),
+        include_str!("asm/smt/utils.asm"),
         include_str!("asm/rlp/decode.asm"),
         include_str!("asm/rlp/encode.asm"),
         include_str!("asm/rlp/encode_rlp_scalar.asm"),