feat: (levm) Flow operations (#553)

**Motivation**

<!-- Why does this pull request exist? What are its goals? -->

**Description**

<!-- A clear and concise general description of the changes this PR
introduces -->

<!-- Link to issues: Resolves #111, Resolves #222 -->

Closes #490
Closes #491 
Closes #489 
Closes #494

---------

Co-authored-by: Ivan Litteri <67517699+ilitteri@users.noreply.github.com>

crates/levm/src/call_frame.rs

@@ -16,7 +16,7 @@ pub struct CallFrame {
 
 impl CallFrame {
     pub fn next_opcode(&mut self) -> Option<Opcode> {
-        let opcode = self.bytecode.get(self.pc).copied().map(Opcode::from);
+        let opcode = self.opcode_at(self.pc);
         self.increment_pc();
         opcode
     }
@@ -32,4 +32,22 @@ impl CallFrame {
     pub fn pc(&self) -> usize {
         self.pc
     }
+
+    pub fn jump(&mut self, jump_address: U256) {
+        if !self.valid_jump(jump_address) {
+            // Should be a halt when we implement it
+            panic!("Invalid jump");
+        }
+        self.pc = jump_address.as_usize() + 1;
+    }
+
+    fn valid_jump(&self, jump_address: U256) -> bool {
+        self.opcode_at(jump_address.as_usize())
+            .map(|opcode| opcode.eq(&Opcode::JUMPDEST))
+            .is_some_and(|is_jumpdest| is_jumpdest)
+    }
+
+    fn opcode_at(&self, offset: usize) -> Option<Opcode> {
+        self.bytecode.get(offset).copied().map(Opcode::from)
+    }
 }

crates/levm/src/opcodes.rs

@@ -70,12 +70,12 @@ pub enum Opcode {
     MSTORE8 = 0x53,
     // SLOAD = 0x54,
     // SSTORE = 0x55,
-    // JUMP = 0x56,
-    // JUMPI = 0x57,
-    // PC = 0x58,
+    JUMP = 0x56,
+    JUMPI = 0x57,
+    PC = 0x58,
     MSIZE = 0x59,
     // GAS = 0x5A,
-    // JUMPDEST = 0x5B,
+    JUMPDEST = 0x5B,
     // TLOAD = 0x5C,
     // TSTORE = 0x5D,
     MCOPY = 0x5E,
@@ -200,6 +200,10 @@ impl From<u8> for Opcode {
             0x14 => Opcode::EQ,
             0x15 => Opcode::ISZERO,
             0x20 => Opcode::KECCAK256,
+            0x56 => Opcode::JUMP,
+            0x57 => Opcode::JUMPI,
+            0x58 => Opcode::PC,
+            0x5B => Opcode::JUMPDEST,
             x if x == Opcode::PUSH0 as u8 => Opcode::PUSH0,
             x if x == Opcode::PUSH1 as u8 => Opcode::PUSH1,
             x if x == Opcode::PUSH2 as u8 => Opcode::PUSH2,

crates/levm/src/operations.rs

@@ -73,12 +73,12 @@ pub enum Operation {
     Mstore8,
     // Sload,
     // Sstore,
-    // Jump,
-    // Jumpi,
-    // PC { pc: usize },
+    Jump,
+    Jumpi,
+    PC,
     Msize,
     // Gas,
-    // Jumpdest { pc: usize },
+    Jumpdest,
     // Tload,
     // Tstore,
     Mcopy,
@@ -172,12 +172,12 @@ impl Operation {
             Operation::Mstore8 => Bytes::copy_from_slice(&[Opcode::MSTORE8 as u8]),
             // Operation::Sload => Bytes::copy_from_slice(&[Opcode::SLOAD as u8]),
             // Operation::Sstore => Bytes::copy_from_slice(&[Opcode::SSTORE as u8]),
-            // Operation::Jump => Bytes::copy_from_slice(&[Opcode::JUMP as u8]),
-            // Operation::Jumpi => Bytes::copy_from_slice(&[Opcode::JUMPI as u8]),
-            // Operation::PC { pc: _ } => Bytes::copy_from_slice(&[Opcode::PC as u8]),
+            Operation::Jump => Bytes::copy_from_slice(&[Opcode::JUMP as u8]),
+            Operation::Jumpi => Bytes::copy_from_slice(&[Opcode::JUMPI as u8]),
+            Operation::PC => Bytes::copy_from_slice(&[Opcode::PC as u8]),
             Operation::Msize => Bytes::copy_from_slice(&[Opcode::MSIZE as u8]),
             // Operation::Gas => Bytes::copy_from_slice(&[Opcode::GAS as u8]),
-            // Operation::Jumpdest { pc: _ } => Bytes::copy_from_slice(&[Opcode::JUMPDEST as u8]),
+            Operation::Jumpdest => Bytes::copy_from_slice(&[Opcode::JUMPDEST as u8]),
             // Operation::Tload => Bytes::copy_from_slice(&[Opcode::TLOAD as u8]),
             // Operation::Tstore => Bytes::copy_from_slice(&[Opcode::TSTORE as u8]),
             Operation::Mcopy => Bytes::copy_from_slice(&[Opcode::MCOPY as u8]),

crates/levm/src/vm.rs

@@ -285,6 +285,25 @@ impl VM {
                         .stack
                         .push(U256::from_big_endian(&result));
                 }
+                Opcode::JUMP => {
+                    let jump_address = current_call_frame.stack.pop().unwrap();
+                    current_call_frame.jump(jump_address);
+                }
+                Opcode::JUMPI => {
+                    let jump_address = current_call_frame.stack.pop().unwrap();
+                    let condition = current_call_frame.stack.pop().unwrap();
+                    if condition != U256::zero() {
+                        current_call_frame.jump(jump_address);
+                    }
+                }
+                Opcode::JUMPDEST => {
+                    // just consume some gas, jumptable written at the start
+                }
+                Opcode::PC => {
+                    current_call_frame
+                        .stack
+                        .push(U256::from(current_call_frame.pc - 1));
+                }
                 Opcode::PUSH0 => {
                     current_call_frame.stack.push(U256::zero());
                 }

crates/levm/tests/tests.rs

@@ -14,7 +14,7 @@ pub fn new_vm_with_ops(operations: &[Operation]) -> VM {
 }
 
 #[test]
-fn test() {
+fn add_op() {
     let mut vm = new_vm_with_ops(&[
         Operation::Push32(U256::one()),
         Operation::Push32(U256::zero()),
@@ -26,8 +26,6 @@ fn test() {
 
     assert!(vm.current_call_frame().stack.pop().unwrap() == U256::one());
     assert!(vm.current_call_frame().pc() == 68);
-
-    println!("{vm:?}");
 }
 
 #[test]
@@ -1166,3 +1164,119 @@ fn pop_on_empty_stack() {
 
     assert!(vm.current_call_frame().stack.pop().unwrap() == U256::one());
 }
+
+#[test]
+fn pc_op() {
+    let operations = [Operation::PC, Operation::Stop];
+    let mut vm = new_vm_with_ops(&operations);
+
+    vm.execute();
+
+    assert!(vm.current_call_frame().stack.pop().unwrap() == U256::from(0));
+}
+
+#[test]
+fn pc_op_with_push_offset() {
+    let operations = [
+        Operation::Push32(U256::one()),
+        Operation::PC,
+        Operation::Stop,
+    ];
+
+    let mut vm = new_vm_with_ops(&operations);
+
+    vm.execute();
+
+    assert!(vm.current_call_frame().stack.pop().unwrap() == U256::from(33));
+}
+
+#[test]
+fn jump_op() {
+    let operations = [
+        Operation::Push32(U256::from(35)),
+        Operation::Jump,
+        Operation::Stop, // should skip this one
+        Operation::Jumpdest,
+        Operation::Push32(U256::from(10)),
+        Operation::Stop,
+    ];
+
+    let mut vm = new_vm_with_ops(&operations);
+
+    vm.execute();
+
+    assert!(vm.current_call_frame().stack.pop().unwrap() == U256::from(10));
+    assert_eq!(vm.current_call_frame().pc(), 70);
+}
+
+#[test]
+#[should_panic]
+fn jump_not_jumpdest_position() {
+    let operations = [
+        Operation::Push32(U256::from(36)),
+        Operation::Jump,
+        Operation::Stop,
+        Operation::Push32(U256::from(10)),
+        Operation::Stop,
+    ];
+
+    let mut vm = new_vm_with_ops(&operations);
+
+    vm.execute();
+    assert_eq!(vm.current_call_frame().pc, 35);
+}
+
+#[test]
+#[should_panic]
+fn jump_position_bigger_than_program_bytecode_size() {
+    let operations = [
+        Operation::Push32(U256::from(5000)),
+        Operation::Jump,
+        Operation::Stop,
+        Operation::Push32(U256::from(10)),
+        Operation::Stop,
+    ];
+
+    let mut vm = new_vm_with_ops(&operations);
+
+    vm.execute();
+    assert_eq!(vm.current_call_frame().pc(), 35);
+}
+
+#[test]
+fn jumpi_not_zero() {
+    let operations = [
+        Operation::Push32(U256::one()),
+        Operation::Push32(U256::from(68)),
+        Operation::Jumpi,
+        Operation::Stop, // should skip this one
+        Operation::Jumpdest,
+        Operation::Push32(U256::from(10)),
+        Operation::Stop,
+    ];
+    let mut vm = new_vm_with_ops(&operations);
+
+    vm.execute();
+
+    assert!(vm.current_call_frame().stack.pop().unwrap() == U256::from(10));
+}
+
+#[test]
+fn jumpi_for_zero() {
+    let operations = [
+        Operation::Push32(U256::from(100)),
+        Operation::Push32(U256::zero()),
+        Operation::Push32(U256::from(100)),
+        Operation::Jumpi,
+        Operation::Stop,
+        Operation::Jumpdest,
+        Operation::Push32(U256::from(10)),
+        Operation::Stop,
+    ];
+
+    let mut vm = new_vm_with_ops(&operations);
+
+    vm.execute();
+
+    assert!(vm.current_call_frame().stack.pop().unwrap() == U256::from(100));
+}