use byteorder::ByteOrder;
use nomos_runtime::ExecutionMessage;
use std::boxed::Box;
use std::collections::HashMap;
use std::str;
use wasmer_runtime::{func, imports, instantiate, Ctx};
pub struct VM {
code: Vec<u8>,
pub state: HashMap<Vec<u8>, Vec<u8>>,
}
pub struct SharedContext {
pub state: HashMap<Vec<u8>, Vec<u8>>,
pub code: Vec<u8>,
}
impl VM {
pub fn new(code: Vec<u8>) -> VM {
let mut state: HashMap<Vec<u8>, Vec<u8>> = HashMap::new();
VM { code, state }
}
pub fn set<K: AsRef<[u8]>, V: AsRef<[u8]>>(&mut self, key: K, value: V) {
self.state
.insert(key.as_ref().to_vec(), value.as_ref().to_vec());
}
pub fn get<K: AsRef<[u8]>>(&self, key: K) -> Option<&Vec<u8>> {
self.state.get(&key.as_ref().to_vec())
}
pub fn call(&mut self, function_name: &str) {
let mut last_state: HashMap<Vec<u8>, Vec<u8>> = self.state.clone();
let mut shared_context = SharedContext {
state: last_state,
code: self.code.clone(),
};
let get_length = |ctx: &mut Ctx, length_result_ptr: u32, key_ptr: u32, key_len: u32| {
let mut shared_context: &mut SharedContext =
unsafe { &mut *(ctx.data as *mut SharedContext) };
let state = &mut shared_context.state;
let memory = ctx.memory(0);
let key_vec: Vec<_> = memory.view()[key_ptr as usize..(key_ptr + key_len) as usize]
.iter()
.map(|cell: &std::cell::Cell<u8>| cell.get())
.collect();
let state_value = state.get(&key_vec);
match state_value {
Some(val) => {
let mut value_length_to_write = [0 as u8; 4];
byteorder::LittleEndian::write_u32(
&mut value_length_to_write,
val.len() as u32,
);
for (byte, cell) in value_length_to_write.iter().zip(
memory.view()[length_result_ptr as usize..(length_result_ptr + 4) as usize]
.iter(),
) {
cell.set(*byte);
}
}
None => (),
};
};
let get_state = |ctx: &mut Ctx, key_ptr: u32, key_len: u32, result_vec_ptr: u32| {
let mut shared_context: &mut SharedContext =
unsafe { &mut *(ctx.data as *mut SharedContext) };
let state = &mut shared_context.state;
let memory = ctx.memory(0);
let key_vec: Vec<_> = memory.view()[key_ptr as usize..(key_ptr + key_len) as usize]
.iter()
.map(|cell: &std::cell::Cell<u8>| cell.get())
.collect();
let state_value = state.get(&key_vec).unwrap();
for (byte, cell) in state_value.iter().zip(
memory.view()
[result_vec_ptr as usize..result_vec_ptr as usize + state_value.len() as usize]
.iter(),
) {
cell.set(*byte);
}
};
let set_state =
|ctx: &mut Ctx, key_ptr: u32, key_len: u32, value_ptr: u32, value_len: u32| {
let mut shared_context: &mut SharedContext =
unsafe { &mut *(ctx.data as *mut SharedContext) };
let state = &mut shared_context.state;
let memory = ctx.memory(0);
let key_vec: Vec<_> = memory.view()[key_ptr as usize..(key_ptr + key_len) as usize]
.iter()
.map(|cell: &std::cell::Cell<u8>| cell.get())
.collect();
let value_vec: Vec<_> = memory.view()
[value_ptr as usize..(value_ptr + value_len) as usize]
.iter()
.map(|cell: &std::cell::Cell<u8>| cell.get())
.collect();
state.insert(key_vec, value_vec);
};
let upgrade_code = |ctx: &mut Ctx, code_ptr: u32, code_len: u32| {
let mut shared_context: &mut SharedContext =
unsafe { &mut *(ctx.data as *mut SharedContext) };
let state = &mut shared_context.state;
let memory = ctx.memory(0);
let code_vec: Vec<_> = memory.view()[code_ptr as usize..(code_ptr + code_len) as usize]
.iter()
.map(|cell: &std::cell::Cell<u8>| cell.get())
.collect();
shared_context.code = code_vec;
};
let execute_code = |ctx: &mut Ctx,
execution_msg_bytes_ptr: u32,
execution_msg_bytes_len: u32| {
let mut shared_context: &mut SharedContext =
unsafe { &mut *(ctx.data as *mut SharedContext) };
let state = &mut shared_context.state;
let memory = ctx.memory(0);
let execution_msg_bytes: Vec<_> = memory.view()[execution_msg_bytes_ptr as usize
..(execution_msg_bytes_ptr + execution_msg_bytes_len) as usize]
.iter()
.map(|cell: &std::cell::Cell<u8>| cell.get())
.collect();
let execution_msg: ExecutionMessage =
bincode::deserialize(&execution_msg_bytes).unwrap();
let code_to_execute = state.get(&execution_msg.code_key);
match code_to_execute {
Some(code) => {
let mut child_vm = VM::new(code.to_vec());
let value_at_store_key = state.get(&execution_msg.store_key);
if let Some(value_bytes) = value_at_store_key {
let child_store_result: Result<
HashMap<Vec<u8>, Vec<u8>>,
Box<bincode::ErrorKind>,
> = bincode::deserialize(&value_bytes);
match child_store_result {
Err(e) => {
return ();
}
Ok(child_store) => {
child_vm.state = child_store;
let execution_result =
child_vm.call(&execution_msg.entry_function_name[..]);
let child_store_bytes =
bincode::serialize(&child_vm.state).unwrap();
state.insert(execution_msg.store_key.to_vec(), child_store_bytes);
}
}
} else {
return ();
};
}
None => {
}
};
};
let dtor = (|_: *mut std::ffi::c_void| {}) as fn(*mut std::ffi::c_void);
let ptr = &mut shared_context as *mut _ as *mut std::ffi::c_void;
let import_object = imports! {
move || {
(ptr, dtor)
},
"env" => {
"print_str" => func!(print_str),
"get_state" => func!(get_state),
"set_state" => func!(set_state),
"get_length" => func!(get_length),
"upgrade_code" => func!(upgrade_code),
"execute_code" => func!(execute_code),
},
};
let mut instance = instantiate(self.code.as_slice(), &import_object).unwrap();
let memory = instance.context_mut().memory(0);
instance
.call(function_name, &[])
.expect("Calling vm method failed");
self.state = shared_context.state.clone();
self.code = shared_context.code.clone();
}
}
fn print_str(ctx: &mut Ctx, ptr: u32, len: u32) {
let memory = ctx.memory(0);
let str_vec: Vec<_> = memory.view()[ptr as usize..(ptr + len) as usize]
.iter()
.map(|cell| cell.get())
.collect();
let string = str::from_utf8(&str_vec).unwrap();
println!("{}", string);
}