Prefetch JUMPDESTs through RPC

Cargo.toml

@@ -35,6 +35,8 @@ alloy = { version = '0.3.0', default-features = false, features = [
   "transport-http",
   "rpc-types-debug",
 ] }
+alloy-primitives = "0.8.0"
+alloy-serde = "0.3.0"
 anyhow = "1.0.86"
 async-stream = "0.3.5"
 axum = "0.7.5"
@@ -50,6 +52,7 @@ criterion = "0.5.1"
 dotenvy = "0.15.7"
 either = "1.12.0"
 enum-as-inner = "0.6.0"
+enumn = "0.1.13"
 env_logger = "0.11.3"
 eth_trie = "0.4.0"
 ethereum-types = "0.14.1"
@@ -86,14 +89,15 @@ serde = "1.0.203"
 serde-big-array = "0.5.1"
 serde_json = "1.0.118"
 serde_path_to_error = "0.1.16"
+serde_with = "3.8.1"
 sha2 = "0.10.8"
 static_assertions = "1.1.0"
 thiserror = "1.0.61"
 tiny-keccak = "2.0.2"
 tokio = { version = "1.38.0", features = ["full"] }
 toml = "0.8.14"
 tower = "0.4"
-tracing = "0.1"
+tracing = { version = "0.1", features = ["attributes"] }
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 u4 = "0.1.0"
 uint = "0.9.5"

evm_arithmetization/benches/fibonacci_25m_gas.rs

@@ -192,6 +192,7 @@ fn prepare_setup() -> anyhow::Result<GenerationInputs<F>> {
             prev_hashes: vec![H256::default(); 256],
             cur_hash: H256::default(),
         },
+        jumpdest_table: Default::default(),
     })
 }
 

evm_arithmetization/src/cpu/kernel/interpreter.rs

@@ -5,11 +5,13 @@
 //! the future execution and generate nondeterministically the corresponding
 //! jumpdest table, before the actual CPU carries on with contract execution.
 
+use core::option::Option::None;
 use std::collections::{BTreeSet, HashMap};
 
 use anyhow::anyhow;
 use ethereum_types::{BigEndianHash, U256};
-use log::Level;
+use keccak_hash::H256;
+use log::{trace, Level};
 use mpt_trie::partial_trie::PartialTrie;
 use plonky2::hash::hash_types::RichField;
 use serde::{Deserialize, Serialize};
@@ -19,8 +21,10 @@ use crate::cpu::columns::CpuColumnsView;
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::generation::debug_inputs;
+use crate::generation::jumpdest::{ContextJumpDests, JumpDestTableProcessed, JumpDestTableWitness};
 use crate::generation::linked_list::LinkedListsPtrs;
 use crate::generation::mpt::{load_linked_lists_and_txn_and_receipt_mpts, TrieRootPtrs};
+use crate::generation::prover_input::get_proofs_and_jumpdests;
 use crate::generation::rlp::all_rlp_prover_inputs_reversed;
 use crate::generation::state::{
     all_ger_prover_inputs, all_withdrawals_prover_inputs_reversed, GenerationState,
@@ -57,6 +61,7 @@ pub(crate) struct Interpreter<F: RichField> {
     /// Counts the number of appearances of each opcode. For debugging purposes.
     #[allow(unused)]
     pub(crate) opcode_count: [usize; 0x100],
+    /// A table of contexts and their reached JUMPDESTs.
     jumpdest_table: HashMap<usize, BTreeSet<usize>>,
     /// `true` if the we are currently carrying out a jumpdest analysis.
     pub(crate) is_jumpdest_analysis: bool,
@@ -72,9 +77,9 @@ pub(crate) struct Interpreter<F: RichField> {
 pub(crate) fn simulate_cpu_and_get_user_jumps<F: RichField>(
     final_label: &str,
     state: &GenerationState<F>,
-) -> Option<HashMap<usize, Vec<usize>>> {
+) -> Option<(JumpDestTableProcessed, JumpDestTableWitness)> {
     match state.jumpdest_table {
-        Some(_) => None,
+        Some(_) => Default::default(),
         None => {
             let halt_pc = KERNEL.global_labels[final_label];
             let initial_context = state.registers.context;
@@ -93,16 +98,15 @@ pub(crate) fn simulate_cpu_and_get_user_jumps<F: RichField>(
 
             let clock = interpreter.get_clock();
 
-            interpreter
+            let jdtw = interpreter
                 .generation_state
-                .set_jumpdest_analysis_inputs(interpreter.jumpdest_table);
+                .set_jumpdest_analysis_inputs(interpreter.jumpdest_table.clone());
 
             log::debug!(
                 "Simulated CPU for jumpdest analysis halted after {:?} cycles.",
                 clock
             );
-
-            interpreter.generation_state.jumpdest_table
+            (interpreter.generation_state.jumpdest_table).map(|x| (x, jdtw))
         }
     }
 }
@@ -115,7 +119,7 @@ pub(crate) struct ExtraSegmentData {
     pub(crate) withdrawal_prover_inputs: Vec<U256>,
     pub(crate) ger_prover_inputs: Vec<U256>,
     pub(crate) trie_root_ptrs: TrieRootPtrs,
-    pub(crate) jumpdest_table: Option<HashMap<usize, Vec<usize>>>,
+    pub(crate) jumpdest_table: Option<JumpDestTableProcessed>,
     pub(crate) access_lists_ptrs: LinkedListsPtrs,
     pub(crate) state_ptrs: LinkedListsPtrs,
     pub(crate) next_txn_index: usize,
@@ -151,6 +155,60 @@ pub(crate) fn set_registers_and_run<F: RichField>(
     interpreter.run()
 }
 
+/// Computes the JUMPDEST proofs for each context.
+///
+/// # Arguments
+///
+/// - `jumpdest_table_rpc`:  The raw table received from RPC.
+/// - `code_db`: The corresponding database of contract code used in the trace.
+pub(crate) fn set_jumpdest_analysis_inputs_rpc(
+    jumpdest_table_rpc: &JumpDestTableWitness,
+    code_map: &HashMap<H256, Vec<u8>>,
+) -> JumpDestTableProcessed {
+    let ctx_proofs = (*jumpdest_table_rpc)
+        .iter()
+        .flat_map(|(code_addr, ctx_jumpdests)| {
+            let code = if code_map.contains_key(code_addr) {
+                &code_map[code_addr]
+            } else {
+                &vec![]
+            };
+            trace!(
+                "code: {:?}, code_addr: {:?}, {:?} <============",
+                &code,
+                &code_addr,
+                code_map.contains_key(code_addr),
+            );
+            trace!("code_map: {:?}", &code_map);
+            prove_context_jumpdests(code, ctx_jumpdests)
+        })
+        .collect();
+    JumpDestTableProcessed::new(ctx_proofs)
+}
+
+/// Orchestrates the proving of all contexts in a specific bytecode.
+///
+/// # Arguments
+///
+/// - `ctx_jumpdests`: Map from `ctx` to its list of offsets to reached
+///   `JUMPDEST`s.
+/// - `code`: The bytecode for the contexts.  This is the same for all contexts.
+fn prove_context_jumpdests(
+    code: &[u8],
+    ctx_jumpdests: &ContextJumpDests,
+) -> HashMap<usize, Vec<usize>> {
+    ctx_jumpdests
+        .0
+        .iter()
+        .map(|(&ctx, jumpdests)| {
+            let proofs = jumpdests.last().map_or(Vec::default(), |&largest_address| {
+                get_proofs_and_jumpdests(code, largest_address, jumpdests.clone())
+            });
+            (ctx, proofs)
+        })
+        .collect()
+}
+
 impl<F: RichField> Interpreter<F> {
     /// Returns an instance of `Interpreter` given `GenerationInputs`, and
     /// assuming we are initializing with the `KERNEL` code.

evm_arithmetization/src/cpu/kernel/tests/add11.rs

@@ -193,6 +193,7 @@ fn test_add11_yml() {
             prev_hashes: vec![H256::default(); 256],
             cur_hash: H256::default(),
         },
+        jumpdest_table: Default::default(),
     };
 
     let initial_stack = vec![];
@@ -370,6 +371,7 @@ fn test_add11_yml_with_exception() {
             prev_hashes: vec![H256::default(); 256],
             cur_hash: H256::default(),
         },
+        jumpdest_table: Default::default(),
     };
 
     let initial_stack = vec![];

evm_arithmetization/src/cpu/kernel/tests/core/jumpdest_analysis.rs

@@ -10,13 +10,14 @@ use plonky2::hash::hash_types::RichField;
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::interpreter::Interpreter;
 use crate::cpu::kernel::opcodes::{get_opcode, get_push_opcode};
+use crate::generation::jumpdest::JumpDestTableProcessed;
 use crate::memory::segments::Segment;
 use crate::witness::memory::MemoryAddress;
 use crate::witness::operation::CONTEXT_SCALING_FACTOR;
 
 impl<F: RichField> Interpreter<F> {
     pub(crate) fn set_jumpdest_analysis_inputs(&mut self, jumps: HashMap<usize, BTreeSet<usize>>) {
-        self.generation_state.set_jumpdest_analysis_inputs(jumps);
+        let _ = self.generation_state.set_jumpdest_analysis_inputs(jumps);
     }
 
     pub(crate) fn get_jumpdest_bit(&self, offset: usize) -> U256 {
@@ -106,7 +107,10 @@ fn test_jumpdest_analysis() -> Result<()> {
         interpreter.generation_state.jumpdest_table,
         // Context 3 has jumpdest 1, 5, 7. All have proof 0 and hence
         // the list [proof_0, jumpdest_0, ... ] is [0, 1, 0, 5, 0, 7, 8, 40]
-        Some(HashMap::from([(3, vec![0, 1, 0, 5, 0, 7, 8, 40])]))
+        Some(JumpDestTableProcessed::new(HashMap::from([(
+            3,
+            vec![0, 1, 0, 5, 0, 7, 8, 40]
+        )])))
     );
 
     // Run jumpdest analysis with context = 3
@@ -123,14 +127,14 @@ fn test_jumpdest_analysis() -> Result<()> {
 
     // We need to manually pop the jumpdest_table and push its value on the top of
     // the stack
-    interpreter
+    (*interpreter
         .generation_state
         .jumpdest_table
         .as_mut()
-        .unwrap()
-        .get_mut(&CONTEXT)
-        .unwrap()
-        .pop();
+        .unwrap())
+    .get_mut(&CONTEXT)
+    .unwrap()
+    .pop();
     interpreter
         .push(41.into())
         .expect("The stack should not overflow");
@@ -175,7 +179,9 @@ fn test_packed_verification() -> Result<()> {
     let mut interpreter: Interpreter<F> =
         Interpreter::new(write_table_if_jumpdest, initial_stack.clone(), None);
     interpreter.set_code(CONTEXT, code.clone());
-    interpreter.generation_state.jumpdest_table = Some(HashMap::from([(3, vec![1, 33])]));
+    interpreter.generation_state.jumpdest_table = Some(JumpDestTableProcessed::new(HashMap::from(
+        [(3, vec![1, 33])],
+    )));
 
     interpreter.run()?;
 
@@ -188,7 +194,9 @@ fn test_packed_verification() -> Result<()> {
         let mut interpreter: Interpreter<F> =
             Interpreter::new(write_table_if_jumpdest, initial_stack.clone(), None);
         interpreter.set_code(CONTEXT, code.clone());
-        interpreter.generation_state.jumpdest_table = Some(HashMap::from([(3, vec![1, 33])]));
+        interpreter.generation_state.jumpdest_table = Some(JumpDestTableProcessed::new(
+            HashMap::from([(3, vec![1, 33])]),
+        ));
 
         assert!(interpreter.run().is_err());
 

evm_arithmetization/src/cpu/kernel/tests/init_exc_stop.rs

@@ -101,6 +101,7 @@ fn test_init_exc_stop() {
             cur_hash: H256::default(),
         },
         ger_data: None,
+        jumpdest_table: Default::default(),
     };
     let initial_stack = vec![];
     let initial_offset = KERNEL.global_labels["init"];