Kernel code to do jumpdest analysis

evm/src/cpu/kernel/aggregator.rs

@@ -16,6 +16,7 @@ pub(crate) fn combined_kernel() -> Kernel {
         include_str!("asm/core/create_addresses.asm"),
         include_str!("asm/core/intrinsic_gas.asm"),
         include_str!("asm/core/invalid.asm"),
+        include_str!("asm/core/jumpdest_analysis.asm"),
         include_str!("asm/core/nonce.asm"),
         include_str!("asm/core/process_txn.asm"),
         include_str!("asm/core/syscall.asm"),

evm/src/cpu/kernel/asm/account_code.asm

@@ -96,8 +96,8 @@ extcodecopy_end:
 // Loads the code at `address` in the `SEGMENT_KERNEL_ACCOUNT_CODE` at the current context and starting at offset 0.
 // Checks that the hash of the loaded code corresponds to the `codehash` in the state trie.
 // Pre stack: address, retdest
-// Post stack: extcodesize(address)
-load_code:
+// Post stack: code_len
+global load_code:
     %stack (address, retdest) -> (extcodehash, address, load_code_ctd, retdest)
     JUMP
 load_code_ctd:

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

@@ -1,10 +1,14 @@
 // Loads some prover-provided contract code into the code segment of memory,
 // then hashes the code and returns the hash.
-
 global bootload_contract:
-    // stack: retdest
+    // stack: address, retdest
+//    %stack (address, retdest) -> (address, after_load_code, retdest)
+//    %jump(load_code)
+    PANIC // TODO
 
-    // TODO
+global bootload_code:
+    // stack: code_len, retdest
+    PANIC // TODO
 
     // stack: code_hash, retdest
     SWAP1

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

@@ -0,0 +1,64 @@
+// Populates @SEGMENT_JUMPDEST_BITS for the given context's code.
+// Pre stack: ctx, code_len, retdest
+// Post stack: (empty)
+global jumpdest_analysis:
+    // stack: ctx, code_len, retdest
+    PUSH 0 // i = 0
+
+loop:
+    // stack: i, ctx, code_len, retdest
+    // Ideally we would break if i >= code_len, but checking i > code_len is
+    // cheaper. It doesn't hurt to over-read by 1, since we'll read 0 which is
+    // a no-op.
+    DUP3 DUP2 GT // i > code_len
+    %jumpi(return)
+
+    // stack: i, ctx, code_len, retdest
+    %stack (i, ctx) -> (ctx, @SEGMENT_CODE, i, i, ctx)
+    MLOAD_GENERAL
+    // stack: opcode, i, ctx, code_len, retdest
+
+    DUP1 %eq_const(0x5b)
+    // stack: opcode == JUMPDEST, opcode, i, ctx, code_len, retdest
+    %jumpi(encountered_jumpdest)
+
+    // stack: opcode, i, ctx, code_len, retdest
+    %code_bytes_to_skip
+    // stack: bytes_to_skip, i, ctx, code_len, retdest
+    ADD
+    %jump(continue)
+
+encountered_jumpdest:
+    // stack: opcode, i, ctx, code_len, retdest
+    POP
+    // stack: i, ctx, code_len, retdest
+    %stack (i, ctx) -> (ctx, @SEGMENT_JUMPDEST_BITS, i, 1, i, ctx)
+    MSTORE_GENERAL
+
+continue:
+    // stack: i, ctx, code_len, retdest
+    %increment
+    %jump(loop)
+
+return:
+    // stack: i, ctx, code_len, retdest
+    %pop3
+    JUMP
+
+// Determines how many bytes to skip, if any, based on the opcode we read.
+// If we read a PUSH<n> opcode, we skip over n bytes, otherwise we skip 0.
+//
+// Note that the range of PUSH opcodes is [0x60, 0x80). I.e. PUSH1 is 0x60
+// and PUSH32 is 0x7f.
+%macro code_bytes_to_skip
+    // stack: opcode
+    %sub_const(0x60)
+    // stack: opcode - 0x60
+    DUP1 %lt_const(0x20)
+    // stack: is_push_opcode, opcode - 0x60
+    SWAP1
+    %increment // n = opcode - 0x60 + 1
+    // stack: n, is_push_opcode
+    MUL
+    // stack: bytes_to_skip
+%endmacro

evm/src/cpu/kernel/interpreter.rs

@@ -202,6 +202,25 @@ impl<'a> Interpreter<'a> {
             rlp.into_iter().map(U256::from).collect();
     }
 
+    pub(crate) fn set_code(&mut self, context: usize, code: Vec<u8>) {
+        assert_ne!(context, 0, "Can't modify kernel code.");
+        while self.memory.context_memory.len() <= context {
+            self.memory
+                .context_memory
+                .push(MemoryContextState::default());
+        }
+        self.memory.context_memory[context].segments[Segment::Code as usize].content =
+            code.into_iter().map(U256::from).collect();
+    }
+
+    pub(crate) fn get_jumpdest_bits(&self, context: usize) -> Vec<bool> {
+        self.memory.context_memory[context].segments[Segment::JumpdestBits as usize]
+            .content
+            .iter()
+            .map(|x| x.bit(0))
+            .collect()
+    }
+
     fn incr(&mut self, n: usize) {
         self.offset += n;
     }

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

@@ -0,0 +1,42 @@
+use anyhow::Result;
+
+use crate::cpu::kernel::aggregator::KERNEL;
+use crate::cpu::kernel::interpreter::Interpreter;
+use crate::cpu::kernel::opcodes::{get_opcode, get_push_opcode};
+
+#[test]
+fn test_jumpdest_analysis() -> Result<()> {
+    let jumpdest_analysis = KERNEL.global_labels["jumpdest_analysis"];
+    const CONTEXT: usize = 3; // arbitrary
+
+    let add = get_opcode("ADD");
+    let push2 = get_push_opcode(2);
+    let jumpdest = get_opcode("JUMPDEST");
+
+    #[rustfmt::skip]
+    let code: Vec<u8> = vec![
+        add,
+        jumpdest,
+        push2,
+        jumpdest, // part of PUSH2
+        jumpdest, // part of PUSH2
+        jumpdest,
+        add,
+        jumpdest,
+    ];
+
+    let expected_jumpdest_bits = vec![false, true, false, false, false, true, false, true];
+
+    // Contract creation transaction.
+    let initial_stack = vec![0xDEADBEEFu32.into(), code.len().into(), CONTEXT.into()];
+    let mut interpreter = Interpreter::new_with_kernel(jumpdest_analysis, initial_stack);
+    interpreter.set_code(CONTEXT, code);
+    interpreter.run()?;
+    assert_eq!(interpreter.stack(), vec![]);
+    assert_eq!(
+        interpreter.get_jumpdest_bits(CONTEXT),
+        expected_jumpdest_bits
+    );
+
+    Ok(())
+}

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

@@ -1,2 +1,3 @@
 mod create_addresses;
 mod intrinsic_gas;
+mod jumpdest_analysis;

evm/src/memory/segments.rs

@@ -38,10 +38,11 @@ pub(crate) enum Segment {
     /// A table of values 2^i for i=0..255 for use with shift
     /// instructions; initialised by `kernel/asm/shift.asm::init_shift_table()`.
     ShiftTable = 16,
+    JumpdestBits = 17,
 }
 
 impl Segment {
-    pub(crate) const COUNT: usize = 17;
+    pub(crate) const COUNT: usize = 18;
 
     pub(crate) fn all() -> [Self; Self::COUNT] {
         [
@@ -62,6 +63,7 @@ impl Segment {
             Self::TrieEncodedChild,
             Self::TrieEncodedChildLen,
             Self::ShiftTable,
+            Self::JumpdestBits,
         ]
     }
 
@@ -85,6 +87,7 @@ impl Segment {
             Segment::TrieEncodedChild => "SEGMENT_TRIE_ENCODED_CHILD",
             Segment::TrieEncodedChildLen => "SEGMENT_TRIE_ENCODED_CHILD_LEN",
             Segment::ShiftTable => "SEGMENT_SHIFT_TABLE",
+            Segment::JumpdestBits => "SEGMENT_JUMPDEST_BITS",
         }
     }
 
@@ -108,6 +111,7 @@ impl Segment {
             Segment::TrieEncodedChild => 256,
             Segment::TrieEncodedChildLen => 6,
             Segment::ShiftTable => 256,
+            Segment::JumpdestBits => 1,
         }
     }
 }