Skip to content

Commit

Permalink
Add builder for VM State struct (#36)
Browse files Browse the repository at this point in the history
* Add builder for VM State struct

* Rename Vm Builder Struct

* Fmt + Clippy

* Use default for builder, remove options

* clippy + lint

* fix tests after merge

---------

Co-authored-by: Fran <[email protected]>
Co-authored-by: Fran <[email protected]>
  • Loading branch information
3 people authored Jun 19, 2024
1 parent 0d1eaab commit 64e0222
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 59 deletions.
100 changes: 63 additions & 37 deletions src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,69 @@ pub struct CallFrame {
// to support in-memory vs on-disk storage, etc.
pub storage: HashMap<U256, U256>,
}

// I'm not really a fan of this, but it saves up time when
// adding new fields to the vm state, and makes it easier
// to setup certain particular state for the tests .
#[derive(Debug, Clone)]
pub struct VMStateBuilder {
pub registers: [U256; 15],
pub flag_lt_of: bool,
pub flag_gt: bool,
pub flag_eq: bool,
pub current_frame: CallFrame,
pub gas_left: u32,
}
impl Default for VMStateBuilder {
fn default() -> Self {
VMStateBuilder {
registers: [U256::zero(); 15],
flag_lt_of: false,
flag_gt: false,
flag_eq: false,
current_frame: CallFrame::new(vec![]),
gas_left: DEFAULT_GAS_LIMIT,
}
}
}
impl VMStateBuilder {
pub fn new() -> VMStateBuilder {
Default::default()
}
pub fn with_registers(mut self, registers: [U256; 15]) -> VMStateBuilder {
self.registers = registers;
self
}
pub fn with_current_frame(mut self, frame: CallFrame) -> VMStateBuilder {
self.current_frame = frame;
self
}
pub fn eq_flag(mut self, eq: bool) -> VMStateBuilder {
self.flag_eq = eq;
self
}
pub fn gt_flag(mut self, gt: bool) -> VMStateBuilder {
self.flag_gt = gt;
self
}
pub fn lt_of_flag(mut self, lt_of: bool) -> VMStateBuilder {
self.flag_lt_of = lt_of;
self
}
pub fn gas_left(mut self, gas_left: u32) -> VMStateBuilder {
self.gas_left = gas_left;
self
}
pub fn build(self) -> VMState {
VMState {
registers: self.registers,
current_frame: self.current_frame,
flag_eq: self.flag_eq,
flag_gt: self.flag_gt,
flag_lt_of: self.flag_lt_of,
gas_left: Saturating(self.gas_left),
}
}
}
#[derive(Debug, Clone)]
pub struct VMState {
// The first register, r0, is actually always zero and not really used.
Expand Down Expand Up @@ -54,42 +116,6 @@ impl VMState {
}
}

pub fn new_with_flag_state(flag_lt_of: bool, flag_eq: bool, flag_gt: bool) -> Self {
let registers = [U256::zero(); 15];
let current_frame = CallFrame::new(vec![]);
Self {
registers,
flag_lt_of,
flag_gt,
flag_eq,
current_frame,
gas_left: Saturating(DEFAULT_GAS_LIMIT),
}
}

pub fn new_with_registers(registers: [U256; 15]) -> Self {
Self {
registers,
flag_lt_of: false,
flag_gt: false,
flag_eq: false,
current_frame: CallFrame::new(vec![]),
gas_left: Saturating(DEFAULT_GAS_LIMIT),
}
}

pub fn new_with_gas(gas_limit: u32) -> Self {
let registers = [U256::zero(); 15];
Self {
registers,
flag_lt_of: false,
flag_gt: false,
flag_eq: false,
current_frame: CallFrame::new(vec![]),
gas_left: Saturating(gas_limit),
}
}

pub fn load_program(&mut self, program_code: Vec<U256>) {
self.current_frame = CallFrame::new(program_code);
}
Expand Down
62 changes: 40 additions & 22 deletions tests/integration_test.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use era_vm::{run_program, run_program_with_custom_state, state::VMState};
use era_vm::{run_program, run_program_with_custom_state, state::VMStateBuilder};
use std::time::{SystemTime, UNIX_EPOCH};
use u256::U256;
const ARTIFACTS_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/program_artifacts");
Expand Down Expand Up @@ -104,39 +104,52 @@ fn test_add_does_not_run_if_eq_is_not_set() {
#[test]
fn test_add_runs_if_eq_is_set() {
let bin_path = make_bin_path_asm("add_conditional_eq");
let vm_with_custom_flags = VMState::new_with_flag_state(false, true, false);
let vm_with_custom_flags = VMStateBuilder::new().eq_flag(true).build();
let (result, _final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert_eq!(result, U256::from_dec_str("10").unwrap());
}

#[test]
fn test_add_does_run_if_lt_is_set() {
let bin_path = make_bin_path_asm("add_conditional_lt");
let vm_with_custom_flags = VMState::new_with_flag_state(true, false, true);
let vm_with_custom_flags = VMStateBuilder::new().lt_of_flag(true).build();
let (result, _final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert_eq!(result, U256::from_dec_str("10").unwrap());
}

#[test]
fn test_add_does_not_run_if_lt_is_not_set() {
let bin_path = make_bin_path_asm("add_conditional_not_lt");
let vm_with_custom_flags = VMState::new_with_flag_state(true, false, true);
let vm_with_custom_flags = VMStateBuilder::new()
.lt_of_flag(true)
.eq_flag(false)
.gt_flag(true)
.build();
// VMState::new_with_flag_state(true, false, true);
let (result, _final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert_eq!(result, U256::from_dec_str("10").unwrap());
}

#[test]
fn test_add_does_run_if_gt_is_set() {
let bin_path = make_bin_path_asm("add_conditional_gt");
let vm_with_custom_flags = VMState::new_with_flag_state(true, false, true);
let vm_with_custom_flags = VMStateBuilder::new()
.lt_of_flag(true)
.eq_flag(false)
.gt_flag(true)
.build();
let (result, _final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert_eq!(result, U256::from_dec_str("20").unwrap());
}

#[test]
fn test_add_does_not_run_if_gt_is_not_set() {
let bin_path = make_bin_path_asm("add_conditional_not_gt");
let vm_with_custom_flags = VMState::new_with_flag_state(false, false, false);
let vm_with_custom_flags = VMStateBuilder::new()
.lt_of_flag(false)
.eq_flag(false)
.gt_flag(false)
.build();
let (result, _final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert_eq!(result, U256::from_dec_str("0").unwrap());
}
Expand All @@ -149,7 +162,7 @@ fn test_add_sets_overflow_flag() {
let mut registers: [U256; 15] = [U256::zero(); 15];
registers[0] = r1;
registers[1] = r2;
let vm_with_custom_flags = VMState::new_with_registers(registers);
let vm_with_custom_flags = VMStateBuilder::new().with_registers(registers).build();
let (_result, final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert!(final_vm_state.flag_lt_of);
}
Expand All @@ -162,7 +175,7 @@ fn test_add_sets_eq_flag() {
let mut registers: [U256; 15] = [U256::zero(); 15];
registers[0] = r1;
registers[1] = r2;
let vm_with_custom_flags = VMState::new_with_registers(registers);
let vm_with_custom_flags = VMStateBuilder::new().with_registers(registers).build();
let (result, final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert!(final_vm_state.flag_eq);
assert!(result.is_zero());
Expand All @@ -176,7 +189,7 @@ fn test_add_sets_gt_flag_keeps_other_flags_clear() {
let mut registers: [U256; 15] = [U256::zero(); 15];
registers[0] = r1;
registers[1] = r2;
let vm_with_custom_flags = VMState::new_with_registers(registers);
let vm_with_custom_flags = VMStateBuilder::new().with_registers(registers).build();
let (result, final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert!(final_vm_state.flag_gt);
assert!(!final_vm_state.flag_eq);
Expand All @@ -198,7 +211,7 @@ fn test_add_does_not_modify_set_flags() {
registers[1] = r2;
registers[2] = r3;
registers[3] = r4;
let vm_with_custom_flags = VMState::new_with_registers(registers);
let vm_with_custom_flags = VMStateBuilder::new().with_registers(registers).build();
let (_result, final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert!(final_vm_state.flag_lt_of);
assert!(final_vm_state.flag_eq);
Expand All @@ -212,7 +225,7 @@ fn test_sub_flags_r1_rs_keeps_other_flags_clear() {
let mut registers: [U256; 15] = [U256::zero(); 15];
registers[0] = r1;
registers[1] = r2;
let vm_with_custom_flags = VMState::new_with_registers(registers);
let vm_with_custom_flags = VMStateBuilder::new().with_registers(registers).build();
let (_result, final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert!(final_vm_state.flag_lt_of);
assert!(!final_vm_state.flag_gt);
Expand All @@ -227,7 +240,7 @@ fn test_sub_sets_eq_flag_keeps_other_flags_clear() {
let mut registers: [U256; 15] = [U256::zero(); 15];
registers[0] = r1;
registers[1] = r2;
let vm_with_custom_flags = VMState::new_with_registers(registers);
let vm_with_custom_flags = VMStateBuilder::new().with_registers(registers).build();
let (_result, final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert!(final_vm_state.flag_eq);
assert!(!final_vm_state.flag_lt_of);
Expand All @@ -242,7 +255,7 @@ fn test_sub_sets_gt_flag_keeps_other_flags_clear() {
let mut registers: [U256; 15] = [U256::zero(); 15];
registers[0] = r1;
registers[1] = r2;
let vm_with_custom_flags = VMState::new_with_registers(registers);
let vm_with_custom_flags = VMStateBuilder::new().with_registers(registers).build();
let (_result, final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert!(final_vm_state.flag_gt);
assert!(!final_vm_state.flag_eq);
Expand Down Expand Up @@ -275,7 +288,7 @@ fn test_mul_big_asm() {
let mut registers: [U256; 15] = [U256::zero(); 15];
registers[0] = r1;
registers[1] = r2;
let vm_with_custom_flags = VMState::new_with_registers(registers);
let vm_with_custom_flags = VMStateBuilder::new().with_registers(registers).build();

let (_, vm) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);

Expand Down Expand Up @@ -309,7 +322,7 @@ fn test_mul_sets_overflow_flag() {
registers[0] = r1;
registers[1] = r2;

let vm_with_custom_flags = VMState::new_with_registers(registers);
let vm_with_custom_flags = VMStateBuilder::new().with_registers(registers).build();
let (_, vm) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert!(vm.flag_lt_of);
}
Expand All @@ -325,7 +338,7 @@ fn test_mul_stack() {
fn test_mul_conditional_gt_set() {
let bin_path = make_bin_path_asm("mul_conditional_gt");

let vm_with_custom_flags = VMState::new_with_flag_state(false, false, true);
let vm_with_custom_flags = VMStateBuilder::new().gt_flag(true).build();
let (result, _) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert_eq!(result, U256::from_dec_str("42").unwrap());
}
Expand All @@ -334,7 +347,7 @@ fn test_mul_conditional_gt_set() {
fn test_mul_conditional_gt_not_set() {
let bin_path = make_bin_path_asm("mul_conditional_gt");

let vm_with_custom_flags = VMState::new_with_flag_state(false, false, false);
let vm_with_custom_flags = VMStateBuilder::new().build();
let (result, _) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
assert_eq!(result, U256::from_dec_str("0").unwrap());
}
Expand Down Expand Up @@ -403,7 +416,7 @@ fn test_div_stack() {
fn test_div_conditional_gt_set() {
let bin_path = make_bin_path_asm("div_conditional_gt");

let vm_with_custom_flags = VMState::new_with_flag_state(false, false, true);
let vm_with_custom_flags = VMStateBuilder::new().gt_flag(true).build();
let (_, vm) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
let quotient_result = vm.get_register(3);
let remainder_result = vm.get_register(4);
Expand All @@ -416,7 +429,7 @@ fn test_div_conditional_gt_set() {
fn test_div_conditional_gt_not_set() {
let bin_path = make_bin_path_asm("div_conditional_gt");

let vm_with_custom_flags = VMState::new_with_flag_state(false, false, false);
let vm_with_custom_flags = VMStateBuilder::new().build();
let (_, vm) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
let quotient_result = vm.get_register(3);
let remainder_result = vm.get_register(4);
Expand All @@ -429,24 +442,29 @@ fn test_div_conditional_gt_not_set() {
#[test]
fn test_more_complex_program_with_conditionals() {
let bin_path = make_bin_path_asm("add_and_sub_with_conditionals");
let vm_with_custom_flags = VMState::new_with_flag_state(false, true, false);
let vm_with_custom_flags = VMStateBuilder::new()
.eq_flag(true)
.gt_flag(false)
.lt_of_flag(false)
.build();
let (result, _final_vm_state) = run_program_with_custom_state(&bin_path, vm_with_custom_flags);
dbg!(_final_vm_state);
assert_eq!(result, U256::from_dec_str("10").unwrap());
}
#[test]
// This test should run out of gas before
// the program can save a number 3 into the storage.
fn test_runs_out_of_gas_and_stops() {
let bin_path = make_bin_path_asm("add_with_costs");
let vm = VMState::new_with_gas(5510);
let vm = VMStateBuilder::new().gas_left(5510).build();
let (result, _) = run_program_with_custom_state(&bin_path, vm);
assert_eq!(result, U256::from_dec_str("0").unwrap());
}

#[test]
fn test_uses_expected_gas() {
let bin_path = make_bin_path_asm("add_with_costs");
let vm = VMState::new_with_gas(5600);
let vm = VMStateBuilder::new().gas_left(5600).build();
let (result, final_vm_state) = run_program_with_custom_state(&bin_path, vm);
assert_eq!(result, U256::from_dec_str("3").unwrap());
assert_eq!(final_vm_state.gas_left(), 0_u32);
Expand Down

0 comments on commit 64e0222

Please sign in to comment.