Implement JumpDest fetching from RPC.

evm_arithmetization/Cargo.toml

@@ -15,6 +15,7 @@ homepage.workspace = true
 keywords.workspace = true
 
 [dependencies]
+__compat_primitive_types = { workspace = true }
 anyhow = { workspace = true }
 bytes = { workspace = true }
 env_logger = { workspace = true }

evm_arithmetization/benches/fibonacci_25m_gas.rs

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

evm_arithmetization/src/cpu/kernel/interpreter.rs

@@ -19,7 +19,9 @@ 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::mpt::{load_linked_lists_and_txn_and_receipt_mpts, TrieRootPtrs};
+use crate::generation::prover_input::{get_proofs_and_jumpdests, CodeDb};
 use crate::generation::rlp::all_rlp_prover_inputs_reversed;
 use crate::generation::state::{
     all_ger_prover_inputs_reversed, all_withdrawals_prover_inputs_reversed, GenerationState,
@@ -56,6 +58,7 @@ pub(crate) struct Interpreter<F: Field> {
     /// 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,
@@ -71,9 +74,9 @@ pub(crate) struct Interpreter<F: Field> {
 pub(crate) fn simulate_cpu_and_get_user_jumps<F: Field>(
     final_label: &str,
     state: &GenerationState<F>,
-) -> Option<HashMap<usize, Vec<usize>>> {
+) -> (Option<JumpDestTableProcessed>, ContextJumpDests) {
     match state.jumpdest_table {
-        Some(_) => None,
+        Some(_) => (None, Default::default()),
         None => {
             let halt_pc = KERNEL.global_labels[final_label];
             let initial_context = state.registers.context;
@@ -94,14 +97,16 @@ pub(crate) fn simulate_cpu_and_get_user_jumps<F: Field>(
 
             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,
+                ContextJumpDests(interpreter.jumpdest_table),
+            )
         }
     }
 }
@@ -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) next_txn_index: usize,
 }
 
@@ -148,6 +153,49 @@ pub(crate) fn set_registers_and_run<F: Field>(
     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_db: &CodeDb,
+) -> JumpDestTableProcessed {
+    let ctx_proofs = jumpdest_table_rpc
+        .0
+        .iter()
+        .flat_map(|(code_addr, ctx_jumpdests)| {
+            prove_context_jumpdests(&code_db[code_addr], ctx_jumpdests)
+        })
+        .collect();
+    JumpDestTableProcessed(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: Field> 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

@@ -195,6 +195,7 @@ fn test_add11_yml() {
             prev_hashes: vec![H256::default(); 256],
             cur_hash: H256::default(),
         },
+        batch_jumpdest_table: None,
     };
 
     let initial_stack = vec![];
@@ -376,6 +377,7 @@ fn test_add11_yml_with_exception() {
             prev_hashes: vec![H256::default(); 256],
             cur_hash: H256::default(),
         },
+        batch_jumpdest_table: None,
     };
 
     let initial_stack = vec![];

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

@@ -7,6 +7,7 @@ use plonky2::field::goldilocks_field::GoldilocksField as F;
 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::witness::operation::CONTEXT_SCALING_FACTOR;
 
 #[test]
@@ -67,7 +68,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(HashMap::from([(
+            3,
+            vec![0, 1, 0, 5, 0, 7, 8, 40]
+        )])))
     );
 
     // Run jumpdest analysis with context = 3
@@ -89,6 +93,7 @@ fn test_jumpdest_analysis() -> Result<()> {
         .jumpdest_table
         .as_mut()
         .unwrap()
+        .0
         .get_mut(&CONTEXT)
         .unwrap()
         .pop();
@@ -136,7 +141,8 @@ 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(HashMap::from([(3, vec![1, 33])])));
 
     interpreter.run()?;
 
@@ -149,7 +155,8 @@ 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(HashMap::from([(3, vec![1, 33])])));
 
         assert!(interpreter.run().is_err());
 

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

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

evm_arithmetization/src/generation/jumpdest.rs

@@ -0,0 +1,64 @@
+use std::{
+    collections::{BTreeSet, HashMap},
+    fmt::Display,
+};
+
+use keccak_hash::H256;
+use serde::{Deserialize, Serialize};
+
+/// Each `CodeAddress` can be called one or more times, each time creating a new
+/// `Context`. Each `Context` will one or more `JumpDests`.
+#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize, Default)]
+pub struct ContextJumpDests(pub HashMap<usize, BTreeSet<usize>>);
+
+/// The result after proving a `JumpDestTableWitness`.
+#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize, Default)]
+pub(crate) struct JumpDestTableProcessed(pub HashMap<usize, Vec<usize>>);
+
+/// Map `CodeAddress -> (Context -> [JumpDests])`
+#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize, Default)]
+pub struct JumpDestTableWitness(pub HashMap<H256, ContextJumpDests>);
+
+impl JumpDestTableWitness {
+    pub fn insert(&mut self, code_hash: &H256, ctx: usize, offset: usize) {
+        self.0.entry(*code_hash).or_default();
+
+        self.0.get_mut(code_hash).unwrap().0.entry(ctx).or_default();
+
+        self.0
+            .get_mut(code_hash)
+            .unwrap()
+            .0
+            .get_mut(&ctx)
+            .unwrap()
+            .insert(offset);
+
+        assert!(self.0.contains_key(code_hash));
+        assert!(self.0[code_hash].0.contains_key(&ctx));
+        assert!(self.0[code_hash].0[&ctx].contains(&offset));
+    }
+}
+
+impl Display for JumpDestTableWitness {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "=== JumpDest table ===")?;
+
+        for (code, ctxtbls) in &self.0 {
+            write!(f, "codehash: {:?}\n{}", code, ctxtbls)?;
+        }
+        Ok(())
+    }
+}
+
+impl Display for ContextJumpDests {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        for (ctx, offsets) in &self.0 {
+            write!(f, "      ctx: {}, offsets: [", ctx)?;
+            for offset in offsets {
+                write!(f, "{:#10x} ", offset)?;
+            }
+            writeln!(f, "]")?;
+        }
+        Ok(())
+    }
+}

evm_arithmetization/src/generation/mod.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 use anyhow::anyhow;
 use ethereum_types::{Address, BigEndianHash, H256, U256};
+use jumpdest::JumpDestTableWitness;
 use keccak_hash::keccak;
 use log::log_enabled;
 use mpt_trie::partial_trie::{HashedPartialTrie, PartialTrie};
@@ -33,6 +34,7 @@ use crate::util::{h2u, u256_to_usize};
 use crate::witness::memory::{MemoryAddress, MemoryChannel, MemoryState};
 use crate::witness::state::RegistersState;
 
+pub mod jumpdest;
 pub(crate) mod linked_list;
 pub mod mpt;
 pub(crate) mod prover_input;
@@ -93,6 +95,10 @@ pub struct GenerationInputs {
     /// The hash of the current block, and a list of the 256 previous block
     /// hashes.
     pub block_hashes: BlockHashes,
+
+    /// A table listing each JUMPDESTs reached in each call context under
+    /// associated code hash.
+    pub batch_jumpdest_table: Option<JumpDestTableWitness>,
 }
 
 /// A lighter version of [`GenerationInputs`], which have been trimmed
@@ -135,6 +141,10 @@ pub struct TrimmedGenerationInputs {
     /// The hash of the current block, and a list of the 256 previous block
     /// hashes.
     pub block_hashes: BlockHashes,
+
+    /// A list of tables listing each JUMPDESTs reached in each call context
+    /// under associated code hash.
+    pub batch_jumpdest_table: Option<JumpDestTableWitness>,
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize, Default)]
@@ -207,6 +217,7 @@ impl GenerationInputs {
             contract_code: self.contract_code.clone(),
             block_metadata: self.block_metadata.clone(),
             block_hashes: self.block_hashes.clone(),
+            batch_jumpdest_table: self.batch_jumpdest_table.clone(),
         }
     }
 }