Fix bugs in jumpdest analysis (#1474)

* Fix simulation for jumpdest analysis

* Fix bugs in jumpdest analysis

* Update evm/src/cpu/kernel/asm/core/jumpdest_analysis.asm

Co-authored-by: Robin Salen <30937548+Nashtare@users.noreply.github.com>

* Address reviews

---------

Co-authored-by: Robin Salen <30937548+Nashtare@users.noreply.github.com>

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

@@ -164,12 +164,10 @@ global write_table_if_jumpdest:
     // stack: (is_1_at_1 or is_0_at_2_or_3|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
     // stack: (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
 
-    // Compute in_range = 
+    // Compute in_range and has_prefix' = 
-    //   - (0xFF|X⁷)³² for the first 15 bytes
+    //   - in_range = (0xFF|X⁷)³²                     and ~has_prefix' = ~has_prefix OR is_0_at_4, for the first 15 bytes
-    //   - (has_prefix => is_0_at_4 |X⁷)³² for the next 15 bytes
+    //   - in_range = (has_prefix => is_0_at_4 |X⁷)³² and ~has_prefix' = ~has_prefix,              for the next 15 bytes
-    //   - (~has_prefix|X⁷)³² for the last byte
+    //   - in_range = (~has_prefix|X⁷)³²              and ~has_prefix' = ~has_prefix,              for the last byte.
-    // Compute also ~has_prefix = ~has_prefix OR is_0_at_4 for all bytes. We don't need to update ~has_prefix
-    // for the second half but it takes less cycles if we do it.
     DUP2 %shl_const(3)
     NOT
     // stack: (is_0_at_4|X⁷)³²,  (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
@@ -177,93 +175,97 @@ global write_table_if_jumpdest:
     PUSH 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00
     AND
     // stack: (is_0_at_4|X⁷)³¹|0⁸,  (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    DUP2
+    DUP1
-    DUP2
+    // pos 0102030405060708091011121314151617181920212223242526272829303132
+    PUSH 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF0000000000000000000000000000000000
+    AND
+    // stack: (is_0_at_4|X⁷)¹⁵|(0⁸)¹⁷, (is_0_at_4|X⁷)³¹|0⁸, (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    DUP3
+    OR
+    // (~has_prefix'|X⁷)³²,  (is_0_at_4|X⁷)³¹|0⁸, (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    SWAP2
     OR
     // pos 0102030405060708091011121314151617181920212223242526272829303132
     PUSH 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF0000000000000000000000000000000000
     OR
-    // stack: (in_range|X⁷)³², (is_0_at_4|X⁷)³²,  (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    // stack: (in_range|X⁷)³², (~has_prefix'|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    SWAP2
-    OR
-    // stack: (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
 
-    // Compute in_range' = in_range AND
+    // Compute in_range' and ~has_prefix as
-    //   - (0xFF|X⁷)³² for bytes in positions 1-7 and 16-23 
+    //   - in_range' = in_range                                     and has_prefix' = ~has_prefix OR is_0_at_5, for bytes in positions 1-7 and 16-23 
-    //   - (has_prefix => is_0_at_5 |X⁷)³² on the rest
+    //   - in_range' = in_range AND (has_prefix => is_0_at_5 |X⁷)³² and has_prefix' = ~has_prefix,              for the rest.
-    // Compute also ~has_prefix = ~has_prefix OR is_0_at_5 for all bytes.
 
     DUP3 %shl_const(4)
     NOT
-    // stack: (is_0_at_5|X⁷)³²,  (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    // stack: (is_0_at_5|X⁷)³²,  (in_range|X⁷)³², (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    DUP2
+    DUP1
-    DUP2
+    // pos 0102030405060708091011121314151617181920212223242526272829303132
+    PUSH 0xFFFFFFFFFFFFFF0000000000000000FFFFFFFFFFFFFFFF000000000000000000
+    AND
+    // stack: (is_0_at_5|X⁷)⁷|(0⁸)⁸|(is_0_at_5|X⁷)⁸|(0⁸)⁸, (is_0_at_5|X⁷)³²,  (in_range|X⁷)³², (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    DUP4
+    OR
+    // stack: (~has_prefix'|X⁷)³², (is_0_at_5|X⁷)³²,  (in_range|X⁷)³², (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    SWAP3
     OR
     // pos 0102030405060708091011121314151617181920212223242526272829303132
     PUSH 0xFFFFFFFFFFFFFF0000000000000000FFFFFFFFFFFFFFFF000000000000000000
     OR
-    // stack: (in_range'|X⁷)³², (is_0_at_5|X⁷)³²,  (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    SWAP2
-    OR
-    // stack: (~has_prefix|X⁷)³², (in_range'|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    SWAP2
     AND
-    SWAP1 
+    // stack: (in_range'|X⁷)³²,  (~has_prefix'|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
 
-    // Compute in_range' = in_range AND
+    // Compute in_range' and ~has_prefix' as
-    //   - (0xFF|X⁷)³² for bytes in positions 1-3, 8-11, 16-19, and 24-27 
+    //   - in_range' = in_range                                     and ~has_prefix' = ~has_prefix OR is_0_at_6, for bytes in positions 1-3, 8-11, 16-19, and 24-27 
-    //   - (has_prefix => is_0_at_6 |X⁷)³² on the rest
+    //   - in_range' = in_range AND (has_prefix => is_0_at_6 |X⁷)³² and ~has_prefix' = has_prefix,               for the rest.
-    // Compute also that ~has_prefix = ~has_prefix OR is_0_at_4 for all bytes.
-
-    // stack: (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
     DUP3 %shl_const(5)
     NOT
-    // stack: (is_0_at_6|X⁷)³²,  (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    // stack: (is_0_at_6|X⁷)³²,  (in_range|X⁷)³², (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    DUP2
+    DUP1
-    DUP2
+    // pos 0102030405060708091011121314151617181920212223242526272829303132
+    PUSH 0xFFFFFF00000000FFFFFFFF00000000FFFFFFFF00000000FFFFFFFF0000000000
+    AND
+    // stack: (is_0_at_6|X⁷)³|(0⁸)⁴|((is_0_at_6|X⁷)⁴|(0⁸)⁴)³, (is_0_at_6|X⁷)³²,  (in_range|X⁷)³², (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    DUP4
+    OR
+    // stack: (~has_prefix'|X⁷)³², (is_0_at_6|X⁷)³²,  (in_range|X⁷)³², (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    SWAP3
     OR
     // pos 0102030405060708091011121314151617181920212223242526272829303132
     PUSH 0xFFFFFF00000000FFFFFFFF00000000FFFFFFFF00000000FFFFFFFF0000000000
     OR
-    // stack: (in_range'|X⁷)³², (is_0_at_6|X⁷)³²,  (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    SWAP2
-    OR
-    // stack: (~has_prefix|X⁷)³², (in_range'|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    SWAP2
     AND
-    SWAP1 
+    // stack: (in_range'|X⁷)³², (~has_prefix'|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
 
-    // Compute in_range' = in_range AND
+    // Compute in_range' and ~has_prefix' as
-    //   - (0xFF|X⁷)³² for bytes in 1, 4-5, 8-9, 12-13, 16-17, 20-21, 24-25, 28-29
+    //   - in_range' = in_range                                     and ~has_prefix' = has_prefix OR is_0_at_7, for bytes in 1, 4-5, 8-9, 12-13, 16-17, 20-21, 24-25, 28-29
-    //   - (has_prefix => is_0_at_7 |X⁷)³² on the rest
+    //   - in_range' = in_range AND (has_prefix => is_0_at_7 |X⁷)³² and ~has_prefix' = ~has_prefix,             for the rest.
-    // Compute also that ~has_prefix = ~has_prefix OR is_0_at_7 for all bytes.
-
-    // stack: (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
     DUP3 %shl_const(6)
     NOT
-    // stack: (is_0_at_7|X⁷)³²,  (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    // stack: (is_0_at_7|X⁷)³², (in_range|X⁷)³², (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    DUP2
+    DUP1
-    DUP2
+    // pos 0102030405060708091011121314151617181920212223242526272829303132
+    PUSH 0xFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF000000
+    AND
+    // stack:  is_0_at_7|X⁷|(0⁸)²|((is_0_at_7|X⁷)²|(0⁸)²)⁷, (is_0_at_7|X⁷)³², (in_range|X⁷)³², (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    DUP4
+    OR
+    // (~has_prefix'|X⁷)³², (is_0_at_7|X⁷)³²,  (in_range|X⁷)³², (~has_prefix|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    SWAP3
     OR
     // pos 0102030405060708091011121314151617181920212223242526272829303132
     PUSH 0xFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF000000
     OR
-    // stack: (in_range'|X⁷)³², (is_0_at_7|X⁷)³²,  (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    SWAP2
-    OR
-    // stack: (~has_prefix|X⁷)³², (in_range'|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
-    SWAP2
     AND
+    // stack: (in_range'|X⁷)³², (~has_prefix'|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+
+    // Compute in_range' as
+    //   - in_range' = in_range,                                    for odd positions
+    //   - in_range' = in_range AND (has_prefix => is_0_at_8 |X⁷)³², for the rest
+
     SWAP1
-
-    // Compute in_range' = in_range AND
-    //   - (0xFF|X⁷)³² for bytes in odd positions
-    //   - (has_prefix => is_0_at_8 |X⁷)³² on the rest
-
     // stack: (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
     DUP3 %shl_const(7)
     NOT
-    // stack: (is_0_at_8|X⁷)³²,  (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
+    // stack: (is_0_at_8|X⁷)³², (~has_prefix|X⁷)³², (in_range|X⁷)³², packed_opcodes, proof_prefix_addr, ctx, jumpdest, retdest
     OR
     // pos 0102030405060708091011121314151617181920212223242526272829303132
     PUSH 0x00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF
@@ -275,14 +277,18 @@ global write_table_if_jumpdest:
     // pos 0102030405060708091011121314151617181920212223242526272829303132
     PUSH 0x8080808080808080808080808080808080808080808080808080808080808080
     AND
-    %jump_neq_const(0x8080808080808080808080808080808080808080808080808080808080808080, return)
+    %jump_neq_const(0x8080808080808080808080808080808080808080808080808080808080808080, return_pop_opcode)
     POP
     %add_const(32)
 
     // check the remaining path
     %jump(verify_path_and_write_jumpdest_table)
+return_pop_opcode:
+    POP
 return:
-    // stack: proof_prefix_addr, jumpdest, ctx, retdest
+    // stack: proof_prefix_addr, ctx, jumpdest, retdest
+    // or
+    // stack: jumpdest, ctx, proof_prefix_addr, retdest
     %pop3
     JUMP
 

evm/src/cpu/kernel/interpreter.rs

@@ -566,6 +566,14 @@ impl<'a> Interpreter<'a> {
                 .contexts
                 .push(MemoryContextState::default());
         }
+        self.generation_state.memory.set(
+            MemoryAddress::new(
+                context,
+                Segment::ContextMetadata,
+                ContextMetadata::CodeSize.unscale(),
+            ),
+            code.len().into(),
+        );
         self.generation_state.memory.contexts[context].segments[Segment::Code.unscale()].content =
             code.into_iter().map(U256::from).collect();
     }
@@ -584,20 +592,8 @@ impl<'a> Interpreter<'a> {
             .collect()
     }
 
-    pub(crate) fn set_jumpdest_bits(&mut self, context: usize, jumpdest_bits: Vec<bool>) {
+    pub(crate) fn set_jumpdest_analysis_inputs(&mut self, jumps: HashMap<usize, BTreeSet<usize>>) {
-        self.generation_state.memory.contexts[context].segments[Segment::JumpdestBits.unscale()]
+        self.generation_state.set_jumpdest_analysis_inputs(jumps);
-            .content = jumpdest_bits.iter().map(|&x| u256_from_bool(x)).collect();
-        self.generation_state
-            .set_proofs_and_jumpdests(HashMap::from([(
-                context,
-                BTreeSet::from_iter(
-                    jumpdest_bits
-                        .into_iter()
-                        .enumerate()
-                        .filter(|&(_, x)| x)
-                        .map(|(i, _)| i),
-                ),
-            )]));
     }
 
     pub(crate) fn incr(&mut self, n: usize) {

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

@@ -1,3 +1,5 @@
+use std::collections::{BTreeSet, HashMap};
+
 use anyhow::Result;
 use ethereum_types::U256;
 
@@ -37,10 +39,84 @@ fn test_jumpdest_analysis() -> Result<()> {
     ];
     let mut interpreter = Interpreter::new_with_kernel(jumpdest_analysis, initial_stack);
     interpreter.set_code(CONTEXT, code);
-    interpreter.set_jumpdest_bits(CONTEXT, jumpdest_bits);
+    interpreter.set_jumpdest_analysis_inputs(HashMap::from([(
+        3,
+        BTreeSet::from_iter(
+            jumpdest_bits
+                .iter()
+                .enumerate()
+                .filter(|&(_, &x)| x)
+                .map(|(i, _)| i),
+        ),
+    )]));
+
+    assert_eq!(
+        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]
+        Some(HashMap::from([(3, vec![0, 1, 0, 5, 0, 7])]))
+    );
 
     interpreter.run()?;
     assert_eq!(interpreter.stack(), vec![]);
 
+    assert_eq!(jumpdest_bits, interpreter.get_jumpdest_bits(3));
+
+    Ok(())
+}
+
+#[test]
+fn test_packed_verification() -> Result<()> {
+    let jumpdest_analysis = KERNEL.global_labels["jumpdest_analysis"];
+    const CONTEXT: usize = 3; // arbitrary
+
+    let add = get_opcode("ADD");
+    let jumpdest = get_opcode("JUMPDEST");
+
+    // The last push(i=0) is 0x5f which is not a valid opcode. However, this
+    // is still meaningful for the test and makes things easier
+    let mut code: Vec<u8> = std::iter::once(add)
+        .chain(
+            (0..=31)
+                .rev()
+                .map(get_push_opcode)
+                .chain(std::iter::once(jumpdest)),
+        )
+        .collect();
+
+    let jumpdest_bits: Vec<bool> = std::iter::repeat(false)
+        .take(33)
+        .chain(std::iter::once(true))
+        .collect();
+
+    // Contract creation transaction.
+    let initial_stack = vec![
+        0xDEADBEEFu32.into(),
+        code.len().into(),
+        U256::from(CONTEXT) << CONTEXT_SCALING_FACTOR,
+    ];
+    let mut interpreter = Interpreter::new_with_kernel(jumpdest_analysis, initial_stack.clone());
+    interpreter.set_code(CONTEXT, code.clone());
+    interpreter.generation_state.jumpdest_table = Some(HashMap::from([(3, vec![1, 33])]));
+
+    interpreter.run()?;
+
+    assert_eq!(jumpdest_bits, interpreter.get_jumpdest_bits(CONTEXT));
+
+    // If we add 1 to each opcode the jumpdest at position 32 is never a valid jumpdest
+    for i in 1..=32 {
+        code[i] += 1;
+        let mut interpreter =
+            Interpreter::new_with_kernel(jumpdest_analysis, initial_stack.clone());
+        interpreter.set_code(CONTEXT, code.clone());
+        interpreter.generation_state.jumpdest_table = Some(HashMap::from([(3, vec![1, 33])]));
+
+        interpreter.run()?;
+
+        assert!(interpreter.get_jumpdest_bits(CONTEXT).is_empty());
+
+        code[i] -= 1;
+    }
+
     Ok(())
 }

evm/src/generation/mod.rs

@@ -348,7 +348,7 @@ fn simulate_cpu_between_labels_and_get_user_jumps<F: Field>(
     final_label: &str,
     state: &mut GenerationState<F>,
 ) -> Option<HashMap<usize, BTreeSet<usize>>> {
-    if state.jumpdest_proofs.is_some() {
+    if state.jumpdest_table.is_some() {
         None
     } else {
         const JUMP_OPCODE: u8 = 0x56;

evm/src/generation/prover_input.rs

@@ -1,5 +1,5 @@
 use std::cmp::min;
-use std::collections::HashMap;
+use std::collections::{BTreeSet, HashMap};
 use std::mem::transmute;
 use std::str::FromStr;
 
@@ -251,38 +251,39 @@ impl<F: Field> GenerationState<F> {
 
     /// Returns the next used jump address.
     fn run_next_jumpdest_table_address(&mut self) -> Result<U256, ProgramError> {
-        let context = self.registers.context;
+        let context = u256_to_usize(stack_peek(self, 0)? >> CONTEXT_SCALING_FACTOR)?;
-        let code_len = u256_to_usize(self.get_code_len()?.into());
+        let code_len = u256_to_usize(self.get_current_code_len()?.into());
 
-        if self.jumpdest_proofs.is_none() {
+        if self.jumpdest_table.is_none() {
-            self.generate_jumpdest_proofs()?;
+            self.generate_jumpdest_table()?;
         }
 
-        let Some(jumpdest_proofs) = &mut self.jumpdest_proofs else {
+        let Some(jumpdest_table) = &mut self.jumpdest_table else {
             return Err(ProgramError::ProverInputError(
                 ProverInputError::InvalidJumpdestSimulation,
             ));
         };
 
-        if let Some(ctx_jumpdest_proofs) = jumpdest_proofs.get_mut(&self.registers.context)
+        if let Some(ctx_jumpdest_table) = jumpdest_table.get_mut(&context)
-            && let Some(next_jumpdest_address) = ctx_jumpdest_proofs.pop()
+            && let Some(next_jumpdest_address) = ctx_jumpdest_table.pop()
         {
             Ok((next_jumpdest_address + 1).into())
         } else {
-            self.jumpdest_proofs = None;
+            self.jumpdest_table = None;
             Ok(U256::zero())
         }
     }
 
     /// Returns the proof for the last jump address.
     fn run_next_jumpdest_table_proof(&mut self) -> Result<U256, ProgramError> {
-        let Some(jumpdest_proofs) = &mut self.jumpdest_proofs else {
+        let context = u256_to_usize(stack_peek(self, 1)? >> CONTEXT_SCALING_FACTOR)?;
+        let Some(jumpdest_table) = &mut self.jumpdest_table else {
             return Err(ProgramError::ProverInputError(
                 ProverInputError::InvalidJumpdestSimulation,
             ));
         };
-        if let Some(ctx_jumpdest_proofs) = jumpdest_proofs.get_mut(&self.registers.context)
+        if let Some(ctx_jumpdest_table) = jumpdest_table.get_mut(&context)
-            && let Some(next_jumpdest_proof) = ctx_jumpdest_proofs.pop()
+            && let Some(next_jumpdest_proof) = ctx_jumpdest_table.pop()
         {
             Ok(next_jumpdest_proof.into())
         } else {
@@ -295,7 +296,7 @@ impl<F: Field> GenerationState<F> {
 
 impl<F: Field> GenerationState<F> {
     /// Simulate the user's code and store all the jump addresses with their respective contexts.
-    fn generate_jumpdest_proofs(&mut self) -> Result<(), ProgramError> {
+    fn generate_jumpdest_table(&mut self) -> Result<(), ProgramError> {
         let checkpoint = self.checkpoint();
         let memory = self.memory.clone();
 
@@ -309,7 +310,7 @@ impl<F: Field> GenerationState<F> {
             "terminate_common",
             self,
         ) else {
-            self.jumpdest_proofs = Some(HashMap::new());
+            self.jumpdest_table = Some(HashMap::new());
             return Ok(());
         };
 
@@ -318,18 +319,18 @@ impl<F: Field> GenerationState<F> {
         self.memory = memory;
 
         // Find proofs for all contexts
-        self.set_proofs_and_jumpdests(jumpdest_table);
+        self.set_jumpdest_analysis_inputs(jumpdest_table);
 
         Ok(())
     }
 
     /// Given a HashMap containing the contexts and the jumpdest addresses, compute their respective proofs,
     /// by calling `get_proofs_and_jumpdests`
-    pub(crate) fn set_proofs_and_jumpdests(
+    pub(crate) fn set_jumpdest_analysis_inputs(
         &mut self,
-        jumpdest_table: HashMap<usize, std::collections::BTreeSet<usize>>,
+        jumpdest_table: HashMap<usize, BTreeSet<usize>>,
     ) {
-        self.jumpdest_proofs = Some(HashMap::from_iter(jumpdest_table.into_iter().map(
+        self.jumpdest_table = Some(HashMap::from_iter(jumpdest_table.into_iter().map(
             |(ctx, jumpdest_table)| {
                 let code = self.get_code(ctx).unwrap();
                 if let Some(&largest_address) = jumpdest_table.last() {
@@ -347,28 +348,42 @@ impl<F: Field> GenerationState<F> {
     }
 
     fn get_code(&self, context: usize) -> Result<Vec<u8>, ProgramError> {
-        let code_len = self.get_code_len()?;
+        let code_len = self.get_code_len(context)?;
         let code = (0..code_len)
             .map(|i| {
-                u256_to_u8(self.memory.get(MemoryAddress::new(
+                u256_to_u8(
-                    self.registers.context,
+                    self.memory
-                    Segment::Code,
+                        .get(MemoryAddress::new(context, Segment::Code, i)),
-                    i,
+                )
-                )))
             })
             .collect::<Result<Vec<u8>, _>>()?;
         Ok(code)
     }
 
-    fn get_code_len(&self) -> Result<usize, ProgramError> {
+    fn set_code_len(&mut self, len: usize) {
+        self.memory.set(
+            MemoryAddress::new(
+                self.registers.context,
+                Segment::ContextMetadata,
+                ContextMetadata::CodeSize.unscale(),
+            ),
+            len.into(),
+        )
+    }
+
+    fn get_code_len(&self, context: usize) -> Result<usize, ProgramError> {
         let code_len = u256_to_usize(self.memory.get(MemoryAddress::new(
-            self.registers.context,
+            context,
             Segment::ContextMetadata,
             ContextMetadata::CodeSize.unscale(),
         )))?;
         Ok(code_len)
     }
 
+    fn get_current_code_len(&self) -> Result<usize, ProgramError> {
+        self.get_code_len(self.registers.context)
+    }
+
     fn set_jumpdest_bits(&mut self, code: &[u8]) {
         const JUMPDEST_OPCODE: u8 = 0x5b;
         for (pos, opcode) in CodeIterator::new(code) {

evm/src/generation/state.rs

@@ -55,7 +55,7 @@ pub(crate) struct GenerationState<F: Field> {
     /// jump destinations with its corresponding "proof". A "proof" for a jump destination is
     /// either 0 or an address i > 32 in the code (not necessarily pointing to an opcode) such that
     /// for every j in [i, i+32] it holds that code[j] < 0x7f - j + i.
-    pub(crate) jumpdest_proofs: Option<HashMap<usize, Vec<usize>>>,
+    pub(crate) jumpdest_table: Option<HashMap<usize, Vec<usize>>>,
 }
 
 impl<F: Field> GenerationState<F> {
@@ -97,7 +97,7 @@ impl<F: Field> GenerationState<F> {
                 txn_root_ptr: 0,
                 receipt_root_ptr: 0,
             },
-            jumpdest_proofs: None,
+            jumpdest_table: None,
         };
         let trie_root_ptrs = state.preinitialize_mpts(&inputs.tries);
 
@@ -189,7 +189,7 @@ impl<F: Field> GenerationState<F> {
                 txn_root_ptr: 0,
                 receipt_root_ptr: 0,
             },
-            jumpdest_proofs: None,
+            jumpdest_table: None,
         }
     }
 }