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"
@@ -46,6 +48,7 @@ ciborium-io = "0.2.2"
 clap = { version = "4.5.7", features = ["derive", "env"] }
 compat = { path = "compat" }
 criterion = "0.5.1"
+derive_more = { version = "1.0.0", features = ["deref", "deref_mut"] }
 dotenvy = "0.15.7"
 either = "1.12.0"
 enum-as-inner = "0.6.0"
@@ -85,6 +88,7 @@ ruint = "1.12.3"
 serde = "1.0.203"
 serde_json = "1.0.118"
 serde_path_to_error = "0.1.16"
+serde_with = "3.8.1"
 serde-big-array = "0.5.1"
 sha2 = "0.10.8"
 static_assertions = "1.1.0"
@@ -93,8 +97,8 @@ thiserror = "1.0.61"
 tiny-keccak = "2.0.2"
 tokio = { version = "1.38.0", features = ["full"] }
 tower = "0.4"
-tracing = "0.1"
-tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+tracing = { version = "0.1", features = ["attributes"] }
+tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
 trybuild = "1.0"
 u4 = "0.1.0"
 uint = "0.9.5"

evm_arithmetization/Cargo.toml

@@ -17,6 +17,7 @@ keywords.workspace = true
 [dependencies]
 anyhow.workspace = true
 bytes.workspace = true
+derive_more.workspace = true
 env_logger.workspace = true
 ethereum-types.workspace = true
 hashbrown.workspace = true

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: None,
     })
 }
 

evm_arithmetization/src/cpu/kernel/interpreter.rs

@@ -5,10 +5,12 @@
 //! 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 keccak_hash::H256;
 use log::Level;
 use mpt_trie::partial_trie::PartialTrie;
 use plonky2::hash::hash_types::RichField;
@@ -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::{Context, 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,
@@ -56,6 +60,7 @@ pub(crate) struct Interpreter<F: RichField> {
     pub(crate) halt_context: Option<usize>,
     /// Counts the number of appearances of each opcode. For debugging purposes.
     pub(crate) opcode_count: HashMap<Operation, usize>,
+    /// A table of call contexts and the JUMPDEST offsets that they jumped to.
     jumpdest_table: HashMap<usize, BTreeSet<usize>>,
     /// `true` if the we are currently carrying out a jumpdest analysis.
     pub(crate) is_jumpdest_analysis: bool,
@@ -71,9 +76,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;
@@ -92,16 +97,16 @@ pub(crate) fn simulate_cpu_and_get_user_jumps<F: RichField>(
 
             let clock = interpreter.get_clock();
 
-            interpreter
+            let (jdtp, jdtw) = interpreter
                 .generation_state
-                .set_jumpdest_analysis_inputs(interpreter.jumpdest_table);
+                .get_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 = Some(jdtp.clone());
+            Some((jdtp, jdtw))
         }
     }
 }
@@ -114,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,
@@ -150,6 +155,57 @@ 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.
+///
+/// # Output
+///
+/// Returns a [`JumpDestTableProccessed`].
+pub(crate) fn get_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![]
+            };
+            prove_context_jumpdests(code, ctx_jumpdests)
+        })
+        .collect();
+    JumpDestTableProcessed::new(ctx_proofs)
+}
+
+/// Orchestrates the proving of all contexts in a specific bytecode.
+///
+/// # Arguments
+///
+/// - `code`: The bytecode for the context `ctx`.
+/// - `ctx`: Map from `ctx` to its list of `JUMPDEST` offsets.
+///
+/// # Outputs
+///
+/// Returns a [`HashMap`] from `ctx` to [`Vec`] of proofs. Each proofs ia a
+/// pair.
+fn prove_context_jumpdests(code: &[u8], ctx: &Context) -> HashMap<usize, Vec<usize>> {
+    ctx.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: None,
     };
 
     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: None,
     };
 
     let initial_stack = vec![];

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

@@ -10,13 +10,15 @@ 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 (jdtp, _jdtw) = self.generation_state.get_jumpdest_analysis_inputs(jumps);
+        self.generation_state.jumpdest_table = Some(jdtp);
     }
 
     pub(crate) fn get_jumpdest_bit(&self, offset: usize) -> U256 {
@@ -106,7 +108,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
@@ -175,7 +180,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 +195,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: None,
     };
     let initial_stack = vec![];
     let initial_offset = KERNEL.global_labels["init"];