Prefetch JUMPDESTs through RPC

Cargo.toml

@@ -34,6 +34,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 @@ clap = { version = "4.5.7", features = ["derive", "env"] }
 alloy-compat = "0.1.0"
 copyvec = "0.2.0"
 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"
@@ -86,6 +89,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"
@@ -94,8 +98,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

@@ -21,6 +21,7 @@ anyhow.workspace = true
 bitvec.workspace = true
 bytes.workspace = true
 copyvec.workspace = true
+derive_more.workspace = true
 either.workspace = true
 env_logger.workspace = true
 ethereum-types.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(),
         },
+        batch_jumpdest_tables: vec![None],
     })
 }
 

evm_arithmetization/src/cpu/kernel/asm/core/jumpdest_analysis.asm

@@ -1,27 +1,27 @@
-// Set @SEGMENT_JUMPDEST_BITS to one between positions [init_pos, final_pos], 
+// Set @SEGMENT_JUMPDEST_BITS to one between positions [init_pos, final_pos],
 // for the given context's code.
 // Pre stack: init_pos, ctx, final_pos, retdest
 // Post stack: (empty)
 global verify_path_and_write_jumpdest_table:
     SWAP2
     DUP2
-    ADD // final_addr
+    ADD // final_addr = final_pos + ctx, i = init_pos
     // stack: final_addr, ctx, i, retdest
     SWAP2
-    ADD // init_addr
+    ADD // init_addr = i + ctx
 loop:
     // stack: i, final_pos, retdest
     DUP2 DUP2 EQ // i == final_pos
     %jumpi(proof_ok)
     DUP2 DUP2 GT // i > final_pos
     %jumpi(proof_not_ok)
 
-     // stack: i, final_pos, retdest
+    // stack: i, final_pos, retdest
     DUP1
     MLOAD_GENERAL // SEGMENT_CODE == 0
     // stack: opcode, i, final_pos, retdest
 
-    DUP1 
+    DUP1
     // Slightly more efficient than `%eq_const(0x5b) ISZERO`
     PUSH 0x5b
     SUB
@@ -141,7 +141,7 @@ global write_table_if_jumpdest:
     // stack: proof_prefix_addr, ctx, jumpdest, retdest
     // If we are here we need to check that the next 32 bytes are not
     // PUSHXX for XX > 32 - n, n in {1, 32}.
-    
+
     %stack
         (proof_prefix_addr, ctx) ->
         (ctx, proof_prefix_addr, 32, proof_prefix_addr, ctx)
@@ -214,11 +214,11 @@ return:
 // non-deterministically guessing the sequence of jumpdest
 // addresses used during program execution within the current context.
 // For each jumpdest address we also non-deterministically guess
-// a proof, which is another address in the code such that 
+// a proof, which is another address in the code such that
 // is_jumpdest doesn't abort, when the proof is at the top of the stack
 // an the jumpdest address below. If that's the case we set the
 // corresponding bit in @SEGMENT_JUMPDEST_BITS to 1.
-// 
+//
 // stack: ctx, code_len, retdest
 // stack: (empty)
 global jumpdest_analysis:
@@ -235,18 +235,19 @@ global jumpdest_analysis_end:
     %pop2
     JUMP
 check_proof:
+    // stack: address + 1, ctx, code_len, retdest
     // stack: address, ctx, code_len, retdest
     DUP3 DUP2 %assert_le
     %decrement
-    // stack: proof, ctx, code_len, retdest
+    // stack: address, ctx, code_len, retdest
     DUP2 SWAP1
     // stack: address, ctx, ctx, code_len, retdest
     // We read the proof
     PROVER_INPUT(jumpdest_table::next_proof)
     // stack: proof, address, ctx, ctx, code_len, retdest
     %write_table_if_jumpdest
     // stack: ctx, code_len, retdest
-    
+
     %jump(jumpdest_analysis)
 
 %macro jumpdest_analysis

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::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,
@@ -54,7 +58,9 @@
     /// The interpreter will halt only if the current context matches
     /// halt_context
     pub(crate) halt_context: Option<usize>,
-    jumpdest_table: HashMap<usize, BTreeSet<usize>>,
+    /// A table of call contexts and the JUMPDEST offsets that they jumped to.
+    // todo
+    jumpdest_table_interpreter: HashMap<usize, BTreeSet<usize>>,
     /// `true` if the we are currently carrying out a jumpdest analysis.
     pub(crate) is_jumpdest_analysis: bool,
     /// Holds the value of the clock: the clock counts the number of operations
@@ -73,10 +79,9 @@
 pub(crate) fn simulate_cpu_and_get_user_jumps<F: RichField>(
     final_label: &str,
     state: &GenerationState<F>,
-) -> Option<HashMap<usize, Vec<usize>>> {
-    match state.jumpdest_table {
-        Some(_) => None,
-        None => {
+) -> Option<(JumpDestTableProcessed, JumpDestTableWitness)> {
+    match state.jumpdest_tables {
+        _ => {
             let halt_pc = KERNEL.global_labels[final_label];
             let initial_context = state.registers.context;
             let mut interpreter = Interpreter::new_with_state_and_halt_condition(
@@ -90,20 +95,23 @@
 
             let _ = interpreter.run();
 
-            log::trace!("jumpdest table = {:?}", interpreter.jumpdest_table);
+            log::trace!(
+                "jumpdest table = {:?}",
+                interpreter.jumpdest_table_interpreter
+            );
 
             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_interpreter.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))
         }
     }
 }
@@ -116,7 +124,8 @@
     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>>>,
+    // todo
+    pub(crate) jumpdest_table: Vec<Option<JumpDestTableProcessed>>,
     pub(crate) access_lists_ptrs: LinkedListsPtrs,
     pub(crate) state_ptrs: LinkedListsPtrs,
     pub(crate) next_txn_index: usize,
@@ -152,6 +161,57 @@
     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>>,
+    prev_max_batch_ctx: usize,
+) -> 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_with_ctx_offset(ctx_proofs, prev_max_batch_ctx)
+}
+
+/// 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.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.
@@ -164,6 +224,7 @@
         debug_inputs(inputs);
 
         let mut result = Self::new(initial_offset, initial_stack, max_cpu_len_log);
+        result.generation_state.jumpdest_tables = vec![None; inputs.batch_jumpdest_tables.len()];
         result.initialize_interpreter_state(inputs);
         result
     }
@@ -182,7 +243,8 @@
             halt_context: None,
             #[cfg(test)]
             opcode_count: HashMap::new(),
-            jumpdest_table: HashMap::new(),
+            // todo
+            jumpdest_table_interpreter: HashMap::new(),
             is_jumpdest_analysis: false,
             clock: 0,
             max_cpu_len_log,
@@ -214,7 +276,8 @@
             halt_context: Some(halt_context),
             #[cfg(test)]
             opcode_count: HashMap::new(),
-            jumpdest_table: HashMap::new(),
+            // check
+            jumpdest_table_interpreter: HashMap::new(),
             is_jumpdest_analysis: true,
             clock: 0,
             max_cpu_len_log,
@@ -473,14 +536,15 @@
             .content
     }
 
+    // what happens here?
     pub(crate) fn add_jumpdest_offset(&mut self, offset: usize) {
         if let Some(jumpdest_table) = self
-            .jumpdest_table
+            .jumpdest_table_interpreter
             .get_mut(&self.generation_state.registers.context)
         {
             jumpdest_table.insert(offset);
         } else {
-            self.jumpdest_table.insert(
+            self.jumpdest_table_interpreter.insert(
                 self.generation_state.registers.context,
                 BTreeSet::from([offset]),
             );

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(),
         },
+        batch_jumpdest_tables: vec![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(),
         },
+        batch_jumpdest_tables: vec![None],
     };
 
     let initial_stack = vec![];