From dc236969b23cf2e1f900cde80edb247e2d4abbe0 Mon Sep 17 00:00:00 2001
From: Linda Guiga <101227802+LindaGuiga@users.noreply.github.com>
Date: Thu, 25 Apr 2024 14:59:03 +0100
Subject: [PATCH] Add dummy segment to the left (#185)

* Add kernel code to MemBefore

* Remove kernel code hashing from

* Fix comments

* Add initial memory check to verifier

* Remove dead code

* Address comment

* Add dummy segments to the left.

* Remove comment and make add11_segments_aggreg use a dummy segment

* Change dummy insertion.

* Apply comments

* Fix Clippy
---
 .../src/cpu/kernel/interpreter.rs             |  31 ++++-
 .../src/cpu/kernel/tests/init_exc_stop.rs     | 107 ++++++++++++++++++
 .../src/cpu/kernel/tests/mod.rs               |   1 +
 evm_arithmetization/src/generation/mod.rs     |   9 +-
 evm_arithmetization/src/generation/state.rs   |  15 ++-
 .../src/memory/memory_stark.rs                |  24 +++-
 evm_arithmetization/src/prover.rs             |  31 ++++-
 evm_arithmetization/src/verifier.rs           |   2 +-
 8 files changed, 204 insertions(+), 16 deletions(-)
 create mode 100644 evm_arithmetization/src/cpu/kernel/tests/init_exc_stop.rs

diff --git a/evm_arithmetization/src/cpu/kernel/interpreter.rs b/evm_arithmetization/src/cpu/kernel/interpreter.rs
index 93d5e1203..7fcceb339 100644
--- a/evm_arithmetization/src/cpu/kernel/interpreter.rs
+++ b/evm_arithmetization/src/cpu/kernel/interpreter.rs
@@ -22,7 +22,9 @@ use crate::generation::rlp::all_rlp_prover_inputs_reversed;
 use crate::generation::state::{
     all_withdrawals_prover_inputs_reversed, GenerationState, GenerationStateCheckpoint,
 };
-use crate::generation::{debug_inputs, TrimmedGenerationInputs};
+use crate::generation::{
+    debug_inputs, TrimmedGenerationInputs, NUM_EXTRA_CYCLES_AFTER, NUM_EXTRA_CYCLES_BEFORE,
+};
 use crate::generation::{state::State, GenerationInputs};
 use crate::keccak_sponge::columns::KECCAK_WIDTH_BYTES;
 use crate::keccak_sponge::keccak_sponge_stark::KeccakSpongeOp;
@@ -62,6 +64,8 @@ pub(crate) struct Interpreter<F: Field> {
     pub(crate) clock: usize,
     /// Log of the maximal number of CPU cycles in one segment execution.
     max_cpu_len_log: Option<usize>,
+    /// Indicates whethere this is a dummy run.
+    is_dummy: bool,
 }
 
 /// Structure storing the state of the interpreter's registers.
@@ -290,6 +294,22 @@ impl<F: Field> Interpreter<F> {
         result
     }
 
+    /// Returns an instance of `Interpreter` given `GenerationInputs`, and
+    /// assuming we are initializing with the `KERNEL` code.
+    pub(crate) fn new_dummy_with_generation_inputs(
+        initial_offset: usize,
+        initial_stack: Vec<U256>,
+        inputs: &GenerationInputs,
+    ) -> Self {
+        debug_inputs(inputs);
+
+        let max_cpu_len = Some(NUM_EXTRA_CYCLES_BEFORE + NUM_EXTRA_CYCLES_AFTER);
+        let mut result =
+            Self::new_with_generation_inputs(initial_offset, initial_stack, inputs, max_cpu_len);
+        result.is_dummy = true;
+        result
+    }
+
     pub(crate) fn new(
         initial_offset: usize,
         initial_stack: Vec<U256>,
@@ -307,6 +327,7 @@ impl<F: Field> Interpreter<F> {
             is_jumpdest_analysis: false,
             clock: 0,
             max_cpu_len_log,
+            is_dummy: false,
         };
         interpreter.generation_state.registers.program_counter = initial_offset;
         let initial_stack_len = initial_stack.len();
@@ -338,6 +359,7 @@ impl<F: Field> Interpreter<F> {
             is_jumpdest_analysis: true,
             clock: 0,
             max_cpu_len_log,
+            is_dummy: false,
         }
     }
 
@@ -521,7 +543,7 @@ impl<F: Field> Interpreter<F> {
     }
 
     pub(crate) fn run(&mut self) -> Result<(RegistersState, Option<MemoryState>), anyhow::Error> {
-        let (final_registers, final_mem) = self.run_cpu(self.max_cpu_len_log)?;
+        let (final_registers, final_mem) = self.run_cpu(self.max_cpu_len_log, self.is_dummy)?;
 
         #[cfg(debug_assertions)]
         {
@@ -551,6 +573,11 @@ impl<F: Field> Interpreter<F> {
             .collect::<Vec<_>>()
     }
 
+    /// Returns the max number of CPU cycles.
+    pub(crate) fn get_max_cpu_len_log(&self) -> Option<usize> {
+        self.max_cpu_len_log
+    }
+
     pub(crate) fn get_txn_field(&self, field: NormalizedTxnField) -> U256 {
         // These fields are already scaled by their respective segment.
         self.generation_state.memory.contexts[0].segments[Segment::TxnFields.unscale()]
diff --git a/evm_arithmetization/src/cpu/kernel/tests/init_exc_stop.rs b/evm_arithmetization/src/cpu/kernel/tests/init_exc_stop.rs
new file mode 100644
index 000000000..d63920e19
--- /dev/null
+++ b/evm_arithmetization/src/cpu/kernel/tests/init_exc_stop.rs
@@ -0,0 +1,107 @@
+use std::collections::HashMap;
+
+use ethereum_types::U256;
+use keccak_hash::keccak;
+use keccak_hash::H256;
+use mpt_trie::partial_trie::HashedPartialTrie;
+use mpt_trie::partial_trie::PartialTrie;
+use plonky2::field::goldilocks_field::GoldilocksField as F;
+
+use crate::cpu::kernel::aggregator::KERNEL;
+use crate::cpu::kernel::interpreter::Interpreter;
+use crate::generation::state::State;
+use crate::generation::TrieInputs;
+use crate::generation::NUM_EXTRA_CYCLES_AFTER;
+use crate::generation::NUM_EXTRA_CYCLES_BEFORE;
+use crate::proof::BlockMetadata;
+use crate::proof::TrieRoots;
+use crate::witness::state::RegistersState;
+use crate::{proof::BlockHashes, GenerationInputs, Node};
+
+// Test to check NUM_EXTRA_CYCLES_BEFORE and NUM_EXTRA_CYCLES_AFTER
+#[test]
+fn test_init_exc_stop() {
+    let block_metadata = BlockMetadata {
+        block_number: 1.into(),
+        ..Default::default()
+    };
+
+    let state_trie = HashedPartialTrie::from(Node::Empty);
+    let transactions_trie = HashedPartialTrie::from(Node::Empty);
+    let receipts_trie = HashedPartialTrie::from(Node::Empty);
+    let storage_tries = vec![];
+
+    let mut contract_code = HashMap::new();
+    contract_code.insert(keccak(vec![]), vec![]);
+
+    // No transactions, so no trie roots change.
+    let trie_roots_after = TrieRoots {
+        state_root: state_trie.hash(),
+        transactions_root: transactions_trie.hash(),
+        receipts_root: receipts_trie.hash(),
+    };
+
+    let inputs = GenerationInputs {
+        signed_txn: None,
+        withdrawals: vec![],
+        tries: TrieInputs {
+            state_trie,
+            transactions_trie,
+            receipts_trie,
+            storage_tries,
+        },
+        trie_roots_after,
+        contract_code,
+        checkpoint_state_trie_root: HashedPartialTrie::from(Node::Empty).hash(),
+        block_metadata,
+        txn_number_before: 0.into(),
+        gas_used_before: 0.into(),
+        gas_used_after: 0.into(),
+        block_hashes: BlockHashes {
+            prev_hashes: vec![H256::default(); 256],
+            cur_hash: H256::default(),
+        },
+    };
+    let initial_stack = vec![];
+    let initial_offset = KERNEL.global_labels["init"];
+    let mut interpreter: Interpreter<F> =
+        Interpreter::new_with_generation_inputs(initial_offset, initial_stack, &inputs, None);
+    interpreter.halt_offsets = vec![KERNEL.global_labels["main"]];
+    interpreter.set_is_kernel(true);
+    interpreter.run().expect("Running dummy init failed.");
+
+    assert_eq!(
+        interpreter.get_clock(),
+        NUM_EXTRA_CYCLES_BEFORE,
+        "NUM_EXTRA_CYCLES_BEFORE is set incorrectly."
+    );
+
+    // The registers should not have changed, besides the stack top.
+    let expected_registers = RegistersState {
+        stack_top: interpreter.get_registers().stack_top,
+        check_overflow: interpreter.get_registers().check_overflow,
+        ..RegistersState::new()
+    };
+
+    assert_eq!(
+        interpreter.get_registers(),
+        expected_registers,
+        "Incorrect registers for dummy run."
+    );
+
+    let main_offset = KERNEL.global_labels["main"];
+    let mut interpreter: Interpreter<F> =
+        Interpreter::new_dummy_with_generation_inputs(initial_offset, vec![], &inputs);
+    interpreter.halt_offsets = vec![KERNEL.global_labels["halt_final"]];
+    interpreter.set_is_kernel(true);
+    interpreter.clock = 0;
+    interpreter.run().expect("Running dummy exc_stop failed.");
+
+    // The "-1" comes from the fact that we stop 1 cycle before the max, to allow
+    // for one padding row, which is needed for CPU STARK.
+    assert_eq!(
+        interpreter.get_clock(),
+        NUM_EXTRA_CYCLES_BEFORE + NUM_EXTRA_CYCLES_AFTER - 1,
+        "NUM_EXTRA_CYCLES_AFTER is set incorrectly."
+    );
+}
diff --git a/evm_arithmetization/src/cpu/kernel/tests/mod.rs b/evm_arithmetization/src/cpu/kernel/tests/mod.rs
index 7581eefe7..818a64588 100644
--- a/evm_arithmetization/src/cpu/kernel/tests/mod.rs
+++ b/evm_arithmetization/src/cpu/kernel/tests/mod.rs
@@ -10,6 +10,7 @@ mod core;
 mod ecc;
 mod exp;
 mod hash;
+mod init_exc_stop;
 mod kernel_consistency;
 mod log;
 mod mpt;
diff --git a/evm_arithmetization/src/generation/mod.rs b/evm_arithmetization/src/generation/mod.rs
index 3db4c19f0..ef682cf09 100644
--- a/evm_arithmetization/src/generation/mod.rs
+++ b/evm_arithmetization/src/generation/mod.rs
@@ -43,6 +43,9 @@ use crate::witness::util::mem_write_log;
 /// Number of cycles to go after having reached the halting state. It is
 /// equal to the number of cycles in `exc_stop` + 1.
 pub const NUM_EXTRA_CYCLES_AFTER: usize = 81;
+/// Number of cycles to go before starting the execution: it is the number of
+/// cycles in `init`.
+pub const NUM_EXTRA_CYCLES_BEFORE: usize = 64;
 /// Memory values used to initialize `MemBefore`.
 pub type MemBeforeValues = Vec<(MemoryAddress, U256)>;
 
@@ -387,6 +390,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     // Initialize the state with the one at the end of the
     // previous segment execution, if any.
     let GenerationSegmentData {
+        is_dummy,
         max_cpu_len_log,
         memory,
         registers_before,
@@ -405,7 +409,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     let cpu_res = timed!(
         timing,
         "simulate CPU",
-        simulate_cpu(&mut state, *max_cpu_len_log)
+        simulate_cpu(&mut state, *max_cpu_len_log, *is_dummy)
     );
     if cpu_res.is_err() {
         output_debug_tries(&state)?;
@@ -470,8 +474,9 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
 fn simulate_cpu<F: Field>(
     state: &mut GenerationState<F>,
     max_cpu_len_log: Option<usize>,
+    is_dummy: bool,
 ) -> anyhow::Result<(RegistersState, Option<MemoryState>)> {
-    let (final_registers, mem_after) = state.run_cpu(max_cpu_len_log)?;
+    let (final_registers, mem_after) = state.run_cpu(max_cpu_len_log, is_dummy)?;
 
     let pc = state.registers.program_counter;
     // Setting the values of padding rows.
diff --git a/evm_arithmetization/src/generation/state.rs b/evm_arithmetization/src/generation/state.rs
index dedd4aa37..8e5f95a5a 100644
--- a/evm_arithmetization/src/generation/state.rs
+++ b/evm_arithmetization/src/generation/state.rs
@@ -76,9 +76,13 @@ pub(crate) trait State<F: Field> {
     fn get_context(&self) -> usize;
 
     /// Checks whether we have reached the maximal cpu length.
-    fn at_end_segment(&self, opt_max_cpu_len_log: Option<usize>) -> bool {
-        if let Some(max_cpu_len_log) = opt_max_cpu_len_log {
-            self.get_clock() == (1 << max_cpu_len_log) - NUM_EXTRA_CYCLES_AFTER
+    fn at_end_segment(&self, opt_max_cpu_len: Option<usize>, is_dummy: bool) -> bool {
+        if let Some(max_cpu_len_log) = opt_max_cpu_len {
+            if is_dummy {
+                self.get_clock() == max_cpu_len_log - NUM_EXTRA_CYCLES_AFTER
+            } else {
+                self.get_clock() == (1 << max_cpu_len_log) - NUM_EXTRA_CYCLES_AFTER
+            }
         } else {
             false
         }
@@ -168,6 +172,7 @@ pub(crate) trait State<F: Field> {
     fn run_cpu(
         &mut self,
         max_cpu_len_log: Option<usize>,
+        is_dummy: bool,
     ) -> anyhow::Result<(RegistersState, Option<MemoryState>)>
     where
         Self: Transition<F>,
@@ -183,7 +188,7 @@ pub(crate) trait State<F: Field> {
             let pc = registers.program_counter;
 
             let halt_final = registers.is_kernel && halt_offsets.contains(&pc);
-            if running && (self.at_halt() || self.at_end_segment(max_cpu_len_log)) {
+            if running && (self.at_halt() || self.at_end_segment(max_cpu_len_log, is_dummy)) {
                 running = false;
                 final_registers = registers;
 
@@ -206,7 +211,7 @@ pub(crate) trait State<F: Field> {
                     }
                 } else {
                     if !running {
-                        assert_eq!(self.get_clock() - final_clock, NUM_EXTRA_CYCLES_AFTER - 1);
+                        debug_assert!(self.get_clock() - final_clock == NUM_EXTRA_CYCLES_AFTER - 1);
                     }
                     let final_mem = if let Some(mut mem) = self.get_full_memory() {
                         // Clear memory we will not use again.
diff --git a/evm_arithmetization/src/memory/memory_stark.rs b/evm_arithmetization/src/memory/memory_stark.rs
index d7ba43eff..5cac318f2 100644
--- a/evm_arithmetization/src/memory/memory_stark.rs
+++ b/evm_arithmetization/src/memory/memory_stark.rs
@@ -1,4 +1,5 @@
 use core::marker::PhantomData;
+use std::cmp::max;
 
 use ethereum_types::U256;
 use itertools::Itertools;
@@ -29,7 +30,7 @@ use crate::memory::columns::{
     TIMESTAMP, TIMESTAMP_INV, VIRTUAL_FIRST_CHANGE,
 };
 use crate::memory::VALUE_LIMBS;
-use crate::witness::memory::MemoryOpKind::Read;
+use crate::witness::memory::MemoryOpKind::{self, Read};
 use crate::witness::memory::{MemoryAddress, MemoryOp};
 
 /// Creates the vector of `Columns` corresponding to:
@@ -235,8 +236,6 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
             if (trace_col_vecs[CONTEXT_FIRST_CHANGE][i] == F::ONE)
                 || (trace_col_vecs[SEGMENT_FIRST_CHANGE][i] == F::ONE)
             {
-                // CONTEXT_FIRST_CHANGE and SEGMENT_FIRST_CHANGE should be 0 at the last row, so
-                // the index should never be out of bounds.
                 if i < trace_col_vecs[ADDR_VIRTUAL].len() - 1 {
                     let x_val = trace_col_vecs[ADDR_VIRTUAL][i + 1].to_canonical_u64() as usize;
                     trace_col_vecs[FREQUENCIES][x_val] += F::ONE;
@@ -275,6 +274,25 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
     /// range check, so this method would add two dummy reads to the same
     /// address, say at timestamps 50 and 80.
     fn fill_gaps(memory_ops: &mut Vec<MemoryOp>) {
+        // First, insert padding row at address (0, 0, 0) if the first row doesn't
+        // have a first virtual address at 0.
+        if memory_ops[0].address.virt != 0 {
+            let dummy_addr = MemoryAddress {
+                context: 0,
+                segment: 0,
+                virt: 0,
+            };
+            memory_ops.insert(
+                0,
+                MemoryOp {
+                    filter: false,
+                    timestamp: 1,
+                    address: dummy_addr,
+                    kind: MemoryOpKind::Read,
+                    value: 0.into(),
+                },
+            );
+        }
         let max_rc = memory_ops.len().next_power_of_two() - 1;
         for (mut curr, mut next) in memory_ops.clone().into_iter().tuple_windows() {
             if curr.address.context != next.address.context
diff --git a/evm_arithmetization/src/prover.rs b/evm_arithmetization/src/prover.rs
index 12be79cc9..5e76b0452 100644
--- a/evm_arithmetization/src/prover.rs
+++ b/evm_arithmetization/src/prover.rs
@@ -40,6 +40,8 @@ use crate::witness::state::RegistersState;
 /// Structure holding the data needed to initialize a segment.
 #[derive(Clone, Default, Debug, Serialize, Deserialize)]
 pub struct GenerationSegmentData {
+    /// Indicates whether this corresponds to a dummy segment.
+    pub(crate) is_dummy: bool,
     /// Registers at the start of the segment execution.
     pub(crate) registers_before: RegistersState,
     /// Registers at the end of the segment execution.
@@ -487,6 +489,7 @@ pub fn generate_all_data_segments<F: RichField>(
     );
 
     let mut segment_data = GenerationSegmentData {
+        is_dummy: false,
         registers_before: RegistersState::new(),
         registers_after: RegistersState::new(),
         memory: MemoryState::default(),
@@ -521,6 +524,7 @@ pub fn generate_all_data_segments<F: RichField>(
         all_seg_data.push(segment_data);
 
         segment_data = GenerationSegmentData {
+            is_dummy: false,
             registers_before: updated_registers,
             // `registers_after` will be set correctly at the next iteration.`
             registers_after: updated_registers,
@@ -546,11 +550,32 @@ pub fn generate_all_data_segments<F: RichField>(
 
     // We need at least two segments to prove a segment aggregation.
     if all_seg_data.len() == 1 {
+        let mut interpreter = Interpreter::<F>::new_dummy_with_generation_inputs(
+            KERNEL.global_labels["init"],
+            vec![],
+            inputs,
+        );
+
         let dummy_seg = GenerationSegmentData {
-            registers_before: segment_data.registers_after,
-            ..segment_data
+            is_dummy: true,
+            registers_before: RegistersState::new(),
+            registers_after: RegistersState::new(),
+            max_cpu_len_log: interpreter.get_max_cpu_len_log(),
+            ..all_seg_data[0].clone()
         };
-        all_seg_data.push(dummy_seg);
+        let (updated_registers, mem_after) =
+            set_registers_and_run(dummy_seg.registers_after, &mut interpreter)?;
+        let mut mem_after = mem_after
+            .expect("The interpreter was running, so it should have returned a MemoryState");
+        // During the interpreter initialization, we set the trie data and initialize
+        // `RlpRaw`. But we do not want to pass this information to the first actual
+        // segment in `MemBefore` since the values are not actually accessed in the
+        // dummy generation.
+        mem_after.contexts[0].segments[Segment::RlpRaw.unscale()].content = vec![];
+        mem_after.contexts[0].segments[Segment::TrieData.unscale()].content = vec![];
+        all_seg_data[0].memory = mem_after;
+
+        all_seg_data.insert(0, dummy_seg);
     }
 
     Ok(all_seg_data)
diff --git a/evm_arithmetization/src/verifier.rs b/evm_arithmetization/src/verifier.rs
index de9e6839f..f0056307a 100644
--- a/evm_arithmetization/src/verifier.rs
+++ b/evm_arithmetization/src/verifier.rs
@@ -20,7 +20,7 @@ use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::memory::segments::Segment;
 use crate::memory::VALUE_LIMBS;
-use crate::proof::{AllProof, AllProofChallenges, PublicValues};
+use crate::proof::{AllProof, AllProofChallenges, MemCap, PublicValues};
 use crate::util::h2u;
 
 pub(crate) fn initial_memory_merkle_cap<