Add call frames (#38)

* Fix tests

* Add far_call test

* fix build issues

* restore far call op

* fix tests

* clippy

* cargo fmt

* move call frame to separate file

* clippy, pr comments

* fix build

* format

* pr comments

---------

Co-authored-by: Fran <[email protected]>

programs/far_call.zasm

@@ -0,0 +1,15 @@
+	.text
+	.file	"far_call.zasm"
+	.globl	__entry
+__entry:
+fake_routine:
+
+.func_begin0:
+    sstore r0, r0
+    far_call r1, r1, @fake_routine
+	ret
+
+.func_end0:
+
+	.note.GNU-stack
+	.rodata

src/address_operands.rs

@@ -39,17 +39,17 @@ pub fn address_operands_read(vm: &mut VMState, opcode: &Opcode) -> (TaggedValue,
                     // stack-=[src0 + offset] + src1
                     let (src0, src1) = reg_and_imm_read(vm, opcode);
                     let res = *vm
-                        .current_frame
+                        .current_context_mut()
                         .stack
                         .get_with_offset(src0.value.as_usize());
-                    vm.current_frame.stack.pop(src0.value);
+                    vm.current_context_mut().stack.pop(src0.value);
                     (res, src1)
                 }
                 ImmMemHandlerFlags::UseStackWithOffset => {
                     // stack[src0 + offset] + src1
                     let (src0, src1) = reg_and_imm_read(vm, opcode);
                     let res = vm
-                        .current_frame
+                        .current_context_mut()
                         .stack
                         .get_with_offset(src0.value.as_usize());
 
@@ -58,15 +58,18 @@ pub fn address_operands_read(vm: &mut VMState, opcode: &Opcode) -> (TaggedValue,
                 ImmMemHandlerFlags::UseAbsoluteOnStack => {
                     // stack=[src0 + offset] + src1
                     let (src0, src1) = reg_and_imm_read(vm, opcode);
-                    let res = vm.current_frame.stack.get_absolute(src0.value.as_usize());
+                    let res = vm
+                        .current_context_mut()
+                        .stack
+                        .get_absolute(src0.value.as_usize());
 
                     (*res, src1)
                 }
                 ImmMemHandlerFlags::UseImm16Only => only_imm16_read(vm, opcode),
                 ImmMemHandlerFlags::UseCodePage => {
                     let (src0, src1) = reg_and_imm_read(vm, opcode);
 
-                    let res = vm.current_frame.code_page[src0.value.as_usize()];
+                    let res = vm.current_context_mut().code_page[src0.value.as_usize()];
                     (TaggedValue::new_raw_integer(res), src1)
                 }
             }
@@ -149,20 +152,22 @@ fn address_operands(vm: &mut VMState, opcode: &Opcode, res: (TaggedValue, Option
                 ImmMemHandlerFlags::UseStackWithPushPop => {
                     // stack+=[src0 + offset] + src1
                     let src0 = reg_and_imm_write(vm, OutputOperandPosition::First, opcode);
-                    vm.current_frame.stack.fill_with_zeros(src0.value + 1);
-                    vm.current_frame.stack.store_with_offset(1, res.0);
+                    vm.current_context_mut()
+                        .stack
+                        .fill_with_zeros(src0.value + 1);
+                    vm.current_context_mut().stack.store_with_offset(1, res.0);
                 }
                 ImmMemHandlerFlags::UseStackWithOffset => {
                     // stack[src0 + offset] + src1
                     let src0 = reg_and_imm_write(vm, OutputOperandPosition::First, opcode);
-                    vm.current_frame
+                    vm.current_context_mut()
                         .stack
                         .store_with_offset(src0.value.as_usize(), res.0);
                 }
                 ImmMemHandlerFlags::UseAbsoluteOnStack => {
                     // stack=[src0 + offset] + src1
                     let src0 = reg_and_imm_write(vm, OutputOperandPosition::First, opcode);
-                    vm.current_frame
+                    vm.current_context_mut()
                         .stack
                         .store_absolute(src0.value.as_usize(), res.0);
                 }

src/call_frame.rs

@@ -0,0 +1,44 @@
+use std::{cell::RefCell, num::Saturating, rc::Rc};
+
+use u256::U256;
+
+use crate::{
+    state::Stack,
+    store::{InMemory, Storage},
+};
+
+#[derive(Debug, Clone)]
+pub struct CallFrame {
+    // Max length for this is 1 << 16. Might want to enforce that at some point
+    pub stack: Stack,
+    pub heap: Vec<U256>,
+    // Code memory is word addressable even though instructions are 64 bit wide.
+    // TODO: this is a Vec of opcodes now but it's probably going to switch back to a
+    // Vec<U256> later on, because I believe we have to record memory queries when
+    // fetching code to execute. Check this
+    pub code_page: Vec<U256>,
+    pub pc: u64,
+    /// Storage for the frame using a type that implements the Storage trait.
+    /// The supported types are InMemory and RocksDB storage.
+    pub storage: Rc<RefCell<dyn Storage>>,
+    /// Transient storage should be used for temporary storage within a transaction and then discarded.
+    pub transient_storage: InMemory,
+    pub gas_left: Saturating<u32>,
+}
+impl CallFrame {
+    pub fn new(
+        program_code: Vec<U256>,
+        gas_stipend: u32,
+        storage: Rc<RefCell<dyn Storage>>,
+    ) -> Self {
+        Self {
+            stack: Stack::new(),
+            heap: vec![],
+            code_page: program_code,
+            pc: 0,
+            gas_left: Saturating(gas_stipend),
+            storage,
+            transient_storage: InMemory::default(),
+        }
+    }
+}

src/lib.rs

@@ -1,16 +1,22 @@
 mod address_operands;
+pub mod call_frame;
 mod op_handlers;
 mod opcode;
 mod ptr_operator;
 pub mod state;
-mod store;
+pub mod store;
 pub mod value;
 
-use std::path::PathBuf;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::env;
+use std::rc::Rc;
 
+use call_frame::CallFrame;
 use op_handlers::add::_add;
 use op_handlers::and::_and;
 use op_handlers::div::_div;
+use op_handlers::far_call::far_call;
 use op_handlers::jump::_jump;
 use op_handlers::log::{
     _storage_read, _storage_write, _transient_storage_read, _transient_storage_write,
@@ -24,7 +30,8 @@ use op_handlers::ptr_sub::_ptr_sub;
 use op_handlers::sub::_sub;
 use op_handlers::xor::_xor;
 pub use opcode::Opcode;
-use state::{VMState, VMStateBuilder};
+use state::{VMState, VMStateBuilder, DEFAULT_INITIAL_GAS};
+use store::{InMemory, RocksDB};
 use u256::U256;
 use zkevm_opcode_defs::definitions::synthesize_opcode_decoding_tables;
 use zkevm_opcode_defs::BinopOpcode;
@@ -33,49 +40,68 @@ use zkevm_opcode_defs::LogOpcode;
 use zkevm_opcode_defs::Opcode as Variant;
 use zkevm_opcode_defs::PtrOpcode;
 
-/// Run a vm program with a clean VM state and with in memory storage.
-pub fn run_program_in_memory(bin_path: &str) -> (U256, VMState) {
-    let vm = VMStateBuilder::default().build();
-    run_program_with_custom_state(bin_path, vm)
-}
-
-/// Run a vm program saving the state to a storage file at the given path.
-pub fn run_program_with_storage(bin_path: &str, storage_path: String) -> (U256, VMState) {
-    let vm = VMStateBuilder::default()
-        .with_storage(PathBuf::from(storage_path))
-        .build();
-    run_program_with_custom_state(bin_path, vm)
-}
-
-/// Run a vm program from the given path using a custom state.
-/// Returns the value stored at storage with key 0 and the final vm state.
-pub fn run_program_with_custom_state(bin_path: &str, mut vm: VMState) -> (U256, VMState) {
-    let opcode_table = synthesize_opcode_decoding_tables(11, ISAVersion(2));
-
+pub fn program_from_file(bin_path: &str) -> Vec<U256> {
     let program = std::fs::read(bin_path).unwrap();
     let encoded = String::from_utf8(program.to_vec()).unwrap();
     let bin = hex::decode(&encoded[2..]).unwrap();
 
     let mut program_code = vec![];
-
     for raw_opcode_slice in bin.chunks(32) {
         let mut raw_opcode_bytes: [u8; 32] = [0; 32];
         raw_opcode_bytes.copy_from_slice(&raw_opcode_slice[..32]);
 
         let raw_opcode_u256 = U256::from_big_endian(&raw_opcode_bytes);
         program_code.push(raw_opcode_u256);
     }
+    program_code
+}
 
-    vm.load_program(program_code);
+/// Run a vm program with a clean VM state and with in memory storage.
+pub fn run_program_in_memory(bin_path: &str) -> (U256, VMState) {
+    let mut vm = VMStateBuilder::default().build();
+    let program_code = program_from_file(bin_path);
+    let storage = Rc::new(RefCell::new(InMemory(HashMap::new())));
+    vm.push_frame(program_code, DEFAULT_INITIAL_GAS, storage);
+    run(vm)
+}
+
+/// Run a vm program saving the state to a storage file at the given path.
+pub fn run_program_with_storage(bin_path: &str, storage_path: &str) -> (U256, VMState) {
+    let storage = RocksDB::open(storage_path.into()).unwrap();
+    let storage = Rc::new(RefCell::new(storage));
+    let frame = CallFrame::new(program_from_file(bin_path), DEFAULT_INITIAL_GAS, storage);
+    let vm = VMStateBuilder::default().with_frames(vec![frame]).build();
+    run(vm)
+}
 
+/// Run a vm program with a clean VM state.
+pub fn run_program(bin_path: &str) -> (U256, VMState) {
+    let vm = VMState::new();
+    run_program_with_custom_state(bin_path, vm)
+}
+/// Run a vm program from the given path using a custom state.
+/// Returns the value stored at storage with key 0 and the final vm state.
+pub fn run_program_with_custom_state(bin_path: &str, mut vm: VMState) -> (U256, VMState) {
+    let program = program_from_file(bin_path);
+    let storage = RocksDB::open(env::temp_dir()).unwrap();
+    let storage = Rc::new(RefCell::new(storage));
+    vm.push_frame(program, DEFAULT_INITIAL_GAS, storage);
+    run(vm)
+}
+
+pub fn run(mut vm: VMState) -> (U256, VMState) {
+    let opcode_table = synthesize_opcode_decoding_tables(11, ISAVersion(2));
     loop {
         let opcode = vm.get_opcode(&opcode_table);
 
         if vm.predicate_holds(&opcode.predicate) {
             match opcode.variant {
                 // TODO: Properly handle what happens
                 // when the VM runs out of ergs/gas.
-                _ if vm.gas_left() == 0 => break,
+                _ if vm.running_frames.len() == 1 && vm.current_context().gas_left.0 == 0 => break,
+                _ if vm.current_context().gas_left.0 == 0 => {
+                    vm.pop_frame();
+                }
                 Variant::Invalid(_) => todo!(),
                 Variant::Nop(_) => todo!(),
                 Variant::Add(_) => {
@@ -109,23 +135,26 @@ pub fn run_program_with_custom_state(bin_path: &str, mut vm: VMState) -> (U256,
                     LogOpcode::TransientStorageRead => _transient_storage_read(&mut vm, &opcode),
                     LogOpcode::TransientStorageWrite => _transient_storage_write(&mut vm, &opcode),
                 },
-                Variant::FarCall(_) => todo!(),
+                Variant::FarCall(far_call_variant) => far_call(&mut vm, &far_call_variant),
+                // TODO: This is not how return works. Fix when we have calls between contracts
+                // hooked up.
+                // This is only to keep the context for tests
                 Variant::Ret(_) => {
-                    // TODO: This is not how return works. Fix when we have calls between contracts
-                    // hooked up.
-                    break;
+                    if vm.running_frames.len() > 1 {
+                        vm.pop_frame();
+                    } else {
+                        break;
+                    }
                 }
                 Variant::UMA(_) => todo!(),
             }
         }
-        vm.current_frame.pc += 1;
+        vm.current_context_mut().pc += 1;
         vm.decrease_gas(&opcode);
     }
-    let final_storage_value = vm
-        .current_frame
-        .storage
-        .borrow()
-        .read(&U256::zero())
-        .unwrap();
+    let final_storage_value = match vm.current_context().storage.borrow().read(&U256::zero()) {
+        Ok(value) => value,
+        Err(_) => U256::zero(),
+    };
     (final_storage_value, vm)
 }

src/op_handlers/far_call.rs

@@ -0,0 +1,17 @@
+use std::{cell::RefCell, collections::HashMap, rc::Rc};
+
+use zkevm_opcode_defs::FarCallOpcode;
+
+use crate::{state::VMState, store::InMemory};
+
+pub fn far_call(vm: &mut VMState, opcode: &FarCallOpcode) {
+    match opcode {
+        FarCallOpcode::Normal => {
+            let program_code = vm.current_context().code_page.clone();
+            let stipend = vm.current_context().gas_left;
+            let storage = Rc::new(RefCell::new(InMemory(HashMap::new())));
+            vm.push_frame(program_code, stipend.0 / 32, storage)
+        }
+        _ => todo!(),
+    }
+}

src/op_handlers/jump.rs

@@ -8,6 +8,6 @@ pub fn _jump(vm: &mut VMState, opcode: &Opcode) {
     let (src0, _) = address_operands_read(vm, opcode);
 
     let next_pc = (src0.value & U256::from(u64::MAX)) - 1; // we subtract 1 because the pc will be incremented after this function
-    vm.current_frame.pc = next_pc.as_u64();
+    vm.current_context_mut().pc = next_pc.as_u64();
     address_operands_store(vm, opcode, TaggedValue::new_raw_integer(next_pc));
 }

src/op_handlers/log.rs

@@ -5,7 +5,7 @@ use crate::{state::VMState, store::Storage, value::TaggedValue, Opcode};
 pub fn _storage_write(vm: &mut VMState, opcode: &Opcode) {
     let key = vm.get_register(opcode.src0_index).value;
     let value = vm.get_register(opcode.src1_index).value;
-    vm.current_frame
+    vm.current_context()
         .storage
         .borrow_mut()
         .store(key, value)
@@ -15,7 +15,7 @@ pub fn _storage_write(vm: &mut VMState, opcode: &Opcode) {
 pub fn _storage_read(vm: &mut VMState, opcode: &Opcode) {
     let key = vm.get_register(opcode.src0_index);
     let value = vm
-        .current_frame
+        .current_context()
         .storage
         .borrow()
         .read(&key.value)
@@ -26,7 +26,7 @@ pub fn _storage_read(vm: &mut VMState, opcode: &Opcode) {
 pub fn _transient_storage_write(vm: &mut VMState, opcode: &Opcode) {
     let key = vm.get_register(opcode.src0_index).value;
     let value = vm.get_register(opcode.src1_index).value;
-    vm.current_frame
+    vm.current_context_mut()
         .transient_storage
         .store(key, value)
         .unwrap();
@@ -35,7 +35,7 @@ pub fn _transient_storage_write(vm: &mut VMState, opcode: &Opcode) {
 pub fn _transient_storage_read(vm: &mut VMState, opcode: &Opcode) {
     let key = vm.get_register(opcode.src0_index).value;
     let value = vm
-        .current_frame
+        .current_context()
         .transient_storage
         .read(&key)
         .unwrap_or(U256::zero());

src/op_handlers/mod.rs

@@ -1,6 +1,7 @@
 pub mod add;
 pub mod and;
 pub mod div;
+pub mod far_call;
 pub mod jump;
 pub mod log;
 pub mod mul;