Bring in Poseidon implementation under cdk_erigon feature flag (#577)

* Bring in Poseidon implementation

Co-authored-by: Linda Guiga <[email protected]>
Co-authored-by: Alonso Gonzalez <[email protected]>

* Add Clippy job with cdk_erigon feature activated

* Apply comments

---------

Co-authored-by: Linda Guiga <[email protected]>
Co-authored-by: Alonso Gonzalez <[email protected]>

.env

@@ -8,3 +8,4 @@ LOGIC_CIRCUIT_SIZE=4..21
 MEMORY_CIRCUIT_SIZE=17..24
 MEMORY_BEFORE_CIRCUIT_SIZE=16..23
 MEMORY_AFTER_CIRCUIT_SIZE=7..23
+POSEIDON_CIRCUIT_SIZE=4..25

.github/workflows/ci.yml

@@ -239,5 +239,8 @@ jobs:
       - name: Run cargo clippy
         run: cargo clippy --all-features --all-targets -- -D warnings -A incomplete-features
 
+      - name: Run cargo clippy (with `cdk_erigon` flag)
+        run: cargo clippy --all-features --all-targets --features cdk_erigon -- -D warnings -A incomplete-features
+
       - name: Rustdoc
         run: cargo doc --all

evm_arithmetization/Cargo.toml

@@ -48,6 +48,7 @@ serde-big-array = { workspace = true }
 
 # Local dependencies
 mpt_trie = { workspace = true }
+smt_trie = { workspace = true, optional = true }
 zk_evm_proc_macro = { workspace = true }
 
 [dev-dependencies]
@@ -64,7 +65,7 @@ parallel = [
   "starky/parallel",
 ]
 polygon_pos = []
-cdk_erigon = []
+cdk_erigon = ["smt_trie"]
 
 [[bin]]
 name = "assemble"

evm_arithmetization/src/all_stark.rs

@@ -25,6 +25,11 @@ use crate::logic::LogicStark;
 use crate::memory::memory_stark::MemoryStark;
 use crate::memory::memory_stark::{self, ctl_context_pruning_looking};
 use crate::memory_continuation::memory_continuation_stark::{self, MemoryContinuationStark};
+#[cfg(feature = "cdk_erigon")]
+use crate::poseidon::{
+    columns::POSEIDON_SPONGE_RATE,
+    poseidon_stark::{self, PoseidonStark, FELT_MAX_BYTES},
+};
 
 /// Structure containing all STARKs and the cross-table lookups.
 #[derive(Clone)]
@@ -38,6 +43,8 @@ pub struct AllStark<F: RichField + Extendable<D>, const D: usize> {
     pub(crate) memory_stark: MemoryStark<F, D>,
     pub(crate) mem_before_stark: MemoryContinuationStark<F, D>,
     pub(crate) mem_after_stark: MemoryContinuationStark<F, D>,
+    #[cfg(feature = "cdk_erigon")]
+    pub(crate) poseidon_stark: PoseidonStark<F, D>,
     pub(crate) cross_table_lookups: Vec<CrossTableLookup<F>>,
 }
 
@@ -55,6 +62,8 @@ impl<F: RichField + Extendable<D>, const D: usize> Default for AllStark<F, D> {
             memory_stark: MemoryStark::default(),
             mem_before_stark: MemoryContinuationStark::default(),
             mem_after_stark: MemoryContinuationStark::default(),
+            #[cfg(feature = "cdk_erigon")]
+            poseidon_stark: PoseidonStark::default(),
             cross_table_lookups: all_cross_table_lookups(),
         }
     }
@@ -72,6 +81,8 @@ impl<F: RichField + Extendable<D>, const D: usize> AllStark<F, D> {
             self.memory_stark.num_lookup_helper_columns(config),
             self.mem_before_stark.num_lookup_helper_columns(config),
             self.mem_after_stark.num_lookup_helper_columns(config),
+            #[cfg(feature = "cdk_erigon")]
+            self.poseidon_stark.num_lookup_helper_columns(config),
         ]
     }
 }
@@ -90,6 +101,8 @@ pub enum Table {
     Memory = 6,
     MemBefore = 7,
     MemAfter = 8,
+    #[cfg(feature = "cdk_erigon")]
+    Poseidon = 9,
 }
 
 impl Deref for Table {
@@ -98,12 +111,20 @@ 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, &7, &8][*self as TableIdx]
+        #[cfg(not(feature = "cdk_erigon"))]
+        return [&0, &1, &2, &3, &4, &5, &6, &7, &8][*self as TableIdx];
+
+        #[cfg(feature = "cdk_erigon")]
+        [&0, &1, &2, &3, &4, &5, &6, &7, &8, &9][*self as TableIdx]
     }
 }
 
 /// Number of STARK tables.
-pub(crate) const NUM_TABLES: usize = Table::MemAfter as usize + 1;
+pub const NUM_TABLES: usize = if cfg!(feature = "cdk_erigon") {
+    Table::MemAfter as usize + 2
+} else {
+    Table::MemAfter as usize + 1
+};
 
 impl Table {
     /// Returns all STARK table indices.
@@ -118,6 +139,8 @@ impl Table {
             Self::Memory,
             Self::MemBefore,
             Self::MemAfter,
+            #[cfg(feature = "cdk_erigon")]
+            Self::Poseidon,
         ]
     }
 }
@@ -135,6 +158,12 @@ pub(crate) fn all_cross_table_lookups<F: Field>() -> Vec<CrossTableLookup<F>> {
         ctl_mem_before(),
         ctl_mem_after(),
         ctl_context_pruning(),
+        #[cfg(feature = "cdk_erigon")]
+        ctl_poseidon_simple(),
+        #[cfg(feature = "cdk_erigon")]
+        ctl_poseidon_general_input(),
+        #[cfg(feature = "cdk_erigon")]
+        ctl_poseidon_general_output(),
     ]
 }
 
@@ -307,6 +336,16 @@ fn ctl_memory<F: Field>() -> CrossTableLookup<F> {
         memory_continuation_stark::ctl_data_memory(),
         memory_continuation_stark::ctl_filter(),
     );
+
+    #[cfg(feature = "cdk_erigon")]
+    let poseidon_general_reads = (0..FELT_MAX_BYTES * POSEIDON_SPONGE_RATE).map(|i| {
+        TableWithColumns::new(
+            *Table::Poseidon,
+            poseidon_stark::ctl_looking_memory(i),
+            poseidon_stark::ctl_looking_memory_filter(),
+        )
+    });
+
     let all_lookers = vec![
         cpu_memory_code_read,
         cpu_push_write_ops,
@@ -317,8 +356,12 @@ fn ctl_memory<F: Field>() -> CrossTableLookup<F> {
     .chain(cpu_memory_gp_ops)
     .chain(keccak_sponge_reads)
     .chain(byte_packing_ops)
-    .chain(iter::once(mem_before_ops))
-    .collect();
+    .chain(iter::once(mem_before_ops));
+
+    #[cfg(feature = "cdk_erigon")]
+    let all_lookers = all_lookers.chain(poseidon_general_reads);
+
+    let all_lookers = all_lookers.collect();
     let memory_looked = TableWithColumns::new(
         *Table::Memory,
         memory_stark::ctl_data(),
@@ -368,3 +411,27 @@ fn ctl_mem_after<F: Field>() -> CrossTableLookup<F> {
     );
     CrossTableLookup::new(all_lookers, mem_after_looked)
 }
+
+#[cfg(feature = "cdk_erigon")]
+fn ctl_poseidon_simple<F: Field>() -> CrossTableLookup<F> {
+    CrossTableLookup::new(
+        vec![cpu_stark::ctl_poseidon_simple_op()],
+        poseidon_stark::ctl_looked_simple_op(),
+    )
+}
+
+#[cfg(feature = "cdk_erigon")]
+fn ctl_poseidon_general_input<F: Field>() -> CrossTableLookup<F> {
+    CrossTableLookup::new(
+        vec![cpu_stark::ctl_poseidon_general_input()],
+        poseidon_stark::ctl_looked_general_input(),
+    )
+}
+
+#[cfg(feature = "cdk_erigon")]
+fn ctl_poseidon_general_output<F: Field>() -> CrossTableLookup<F> {
+    CrossTableLookup::new(
+        vec![cpu_stark::ctl_poseidon_general_output()],
+        poseidon_stark::ctl_looked_general_output(),
+    )
+}

evm_arithmetization/src/cpu/columns/ops.rs

@@ -20,6 +20,9 @@ pub(crate) struct OpsColumnsView<T: Copy> {
     pub shift: T,
     /// Combines JUMPDEST and KECCAK_GENERAL flags.
     pub jumpdest_keccak_general: T,
+    /// Combines POSEIDON and POSEIDON_GENERAL flags.
+    #[cfg(feature = "cdk_erigon")]
+    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,8 @@ const KEEPS_CONTEXT: OpsColumnsView<bool> = OpsColumnsView {
     not_pop: true,
     shift: true,
     jumpdest_keccak_general: true,
+    #[cfg(feature = "cdk_erigon")]
+    poseidon: true,
     push_prover_input: true,
     jumps: true,
     pc_push0: true,

evm_arithmetization/src/cpu/control_flow.rs

@@ -8,7 +8,9 @@ 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_INST_LEN: usize = if cfg!(feature = "cdk_erigon") { 13 } else { 12 };
+
+const NATIVE_INSTRUCTIONS: [usize; NATIVE_INST_LEN] = [
     COL_MAP.op.binary_op,
     COL_MAP.op.ternary_op,
     COL_MAP.op.fp254_op,
@@ -17,6 +19,8 @@ const NATIVE_INSTRUCTIONS: [usize; 12] = [
     COL_MAP.op.not_pop,
     COL_MAP.op.shift,
     COL_MAP.op.jumpdest_keccak_general,
+    #[cfg(feature = "cdk_erigon")]
+    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

@@ -462,6 +462,87 @@ pub(crate) fn ctl_filter_set_context<F: Field>() -> Filter<F> {
     )
 }
 
+#[cfg(feature = "cdk_erigon")]
+/// Returns the `TableWithColumns` for the CPU rows calling POSEIDON.
+pub(crate) fn ctl_poseidon_simple_op<F: Field>() -> TableWithColumns<F> {
+    // When executing POSEIDON, the GP memory channels are used as follows:
+    // GP channel 0: stack[-1] = x
+    // GP channel 1: stack[-2] = y
+    // GP channel 2: stack[-3] = z
+    // Such that we can compute `POSEIDON(x || y || z)`.
+    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));
+    TableWithColumns::new(*Table::Cpu, columns, ctl_poseidon_simple_filter())
+}
+
+#[cfg(feature = "cdk_erigon")]
+pub(crate) fn ctl_poseidon_general_input<F: Field>() -> TableWithColumns<F> {
+    // When executing POSEIDON_GENERAL, the GP memory channels are used as follows:
+    // GP channel 0: stack[-1] = addr (context, segment, virt)
+    // GP channel 1: stack[-2] = len
+    let (context, segment, virt) = get_addr(&COL_MAP, 0);
+    let context = Column::single(context);
+    let segment: Column<F> = Column::single(segment);
+    let virt = Column::single(virt);
+    let len = Column::single(COL_MAP.mem_channels[1].value[0]);
+
+    let num_channels = F::from_canonical_usize(NUM_CHANNELS);
+    let timestamp = Column::linear_combination([(COL_MAP.clock, num_channels)]);
+
+    TableWithColumns::new(
+        *Table::Cpu,
+        vec![context, segment, virt, len, timestamp],
+        ctl_poseidon_general_filter(),
+    )
+}
+
+#[cfg(feature = "cdk_erigon")]
+/// CTL filter for the `POSEIDON` operation.
+/// POSEIDON is differentiated from POSEIDON_GENERAL by its first bit set to 0.
+pub(crate) fn ctl_poseidon_simple_filter<F: Field>() -> Filter<F> {
+    Filter::new(
+        vec![(
+            Column::single(COL_MAP.op.poseidon),
+            Column::linear_combination_with_constant([(COL_MAP.opcode_bits[0], -F::ONE)], F::ONE),
+        )],
+        vec![],
+    )
+}
+
+#[cfg(feature = "cdk_erigon")]
+/// CTL filter for the `POSEIDON_GENERAL` operation.
+/// POSEIDON_GENERAL is differentiated from POSEIDON by its first bit set to 1.
+pub(crate) fn ctl_poseidon_general_filter<F: Field>() -> Filter<F> {
+    Filter::new(
+        vec![(
+            Column::single(COL_MAP.op.poseidon),
+            Column::single(COL_MAP.opcode_bits[0]),
+        )],
+        vec![],
+    )
+}
+
+#[cfg(feature = "cdk_erigon")]
+/// Returns the `TableWithColumns` for the CPU rows calling POSEIDON_GENERAL.
+pub(crate) fn ctl_poseidon_general_output<F: Field>() -> TableWithColumns<F> {
+    let mut columns = Vec::new();
+    columns.extend(Column::singles_next_row(COL_MAP.mem_channels[0].value));
+    let num_channels = F::from_canonical_usize(NUM_CHANNELS);
+    columns.push(Column::linear_combination([(COL_MAP.clock, num_channels)]));
+    TableWithColumns::new(*Table::Cpu, columns, ctl_poseidon_general_filter())
+}
+
 /// 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

@@ -7,6 +7,8 @@ use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsume
 
 use crate::cpu::columns::{CpuColumnsView, COL_MAP};
 
+const OPCODES_LEN: usize = if cfg!(feature = "cdk_erigon") { 6 } else { 5 };
+
 /// List of opcode blocks
 /// Each block corresponds to exactly one flag, and each flag corresponds to
 /// exactly one block.  Each block of opcodes:
@@ -29,13 +31,15 @@ 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); OPCODES_LEN] = [
     // (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),
+    #[cfg(feature = "cdk_erigon")]
+    (0x22, 1, 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,10 @@ 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.
+    #[cfg(feature = "cdk_erigon")]
+    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,