use crate::error::Error;
use crate::module::{AddrDetails, GlobalSpec, HeapSpec, Module, ModuleInternal, TableElement};
use libc::c_void;
use lucet_module::owned::{
OwnedExportFunction, OwnedFunctionMetadata, OwnedGlobalSpec, OwnedImportFunction,
OwnedLinearMemorySpec, OwnedModuleData, OwnedSparseData,
};
use lucet_module::{
FunctionHandle, FunctionIndex, FunctionPointer, FunctionSpec, ModuleData, ModuleFeatures,
Signature, TrapSite, UniqueSignatureIndex,
};
use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
#[derive(Default)]
pub struct MockModuleBuilder {
heap_spec: HeapSpec,
sparse_page_data: Vec<Option<Vec<u8>>>,
globals: BTreeMap<usize, OwnedGlobalSpec>,
table_elements: BTreeMap<usize, TableElement>,
export_funcs: HashMap<&'static str, FunctionPointer>,
func_table: HashMap<(u32, u32), FunctionPointer>,
start_func: Option<FunctionPointer>,
function_manifest: Vec<FunctionSpec>,
function_info: Vec<OwnedFunctionMetadata>,
imports: Vec<OwnedImportFunction>,
exports: Vec<OwnedExportFunction>,
signatures: Vec<Signature>,
}
impl MockModuleBuilder {
pub fn new() -> Self {
const DEFAULT_HEAP_SPEC: HeapSpec = HeapSpec {
reserved_size: 4 * 1024 * 1024,
guard_size: 4 * 1024 * 1024,
initial_size: 64 * 1024,
max_size: Some(64 * 1024),
};
MockModuleBuilder::default().with_heap_spec(DEFAULT_HEAP_SPEC)
}
pub fn with_heap_spec(mut self, heap_spec: HeapSpec) -> Self {
self.heap_spec = heap_spec;
self
}
pub fn with_initial_heap(mut self, heap: &[u8]) -> Self {
self.sparse_page_data = heap
.chunks(4096)
.map(|page| {
if page.iter().all(|b| *b == 0) {
None
} else {
let mut page = page.to_vec();
if page.len() < 4096 {
page.resize(4096, 0);
}
Some(page)
}
})
.collect();
self
}
pub fn with_global(mut self, idx: u32, init_val: i64) -> Self {
self.globals
.insert(idx as usize, OwnedGlobalSpec::new_def(init_val, vec![]));
self
}
pub fn with_exported_global(mut self, idx: u32, init_val: i64, export_name: &str) -> Self {
self.globals.insert(
idx as usize,
OwnedGlobalSpec::new_def(init_val, vec![export_name.to_string()]),
);
self
}
pub fn with_import(mut self, idx: u32, import_module: &str, import_field: &str) -> Self {
self.globals.insert(
idx as usize,
OwnedGlobalSpec::new_import(
import_module.to_string(),
import_field.to_string(),
vec![],
),
);
self
}
pub fn with_exported_import(
mut self,
idx: u32,
import_module: &str,
import_field: &str,
export_name: &str,
) -> Self {
self.globals.insert(
idx as usize,
OwnedGlobalSpec::new_import(
import_module.to_string(),
import_field.to_string(),
vec![export_name.to_string()],
),
);
self
}
pub fn with_table_element(mut self, idx: u32, element: &TableElement) -> Self {
self.table_elements.insert(idx as usize, element.clone());
self
}
fn record_sig(&mut self, sig: Signature) -> UniqueSignatureIndex {
let idx = self
.signatures
.iter()
.enumerate()
.find(|(_, v)| *v == &sig)
.map(|(key, _)| key)
.unwrap_or_else(|| {
self.signatures.push(sig);
self.signatures.len() - 1
});
UniqueSignatureIndex::from_u32(idx as u32)
}
pub fn with_export_func(mut self, export: MockExportBuilder) -> Self {
self.export_funcs.insert(export.sym(), export.func());
let sig_idx = self.record_sig(export.sig());
self.function_info.push(OwnedFunctionMetadata {
signature: sig_idx,
name: Some(export.sym().to_string()),
});
self.exports.push(OwnedExportFunction {
fn_idx: FunctionIndex::from_u32(self.function_manifest.len() as u32),
names: vec![export.sym().to_string()],
});
self.function_manifest.push(FunctionSpec::new(
export.func().as_usize() as u64,
export.func_len() as u32,
export.traps().as_ptr() as u64,
export.traps().len() as u64,
));
self
}
pub fn with_exported_import_func(
mut self,
export_name: &'static str,
import_fn_ptr: FunctionPointer,
sig: Signature,
) -> Self {
self.export_funcs.insert(export_name, import_fn_ptr);
let sig_idx = self.record_sig(sig);
self.function_info.push(OwnedFunctionMetadata {
signature: sig_idx,
name: Some(export_name.to_string()),
});
self.exports.push(OwnedExportFunction {
fn_idx: FunctionIndex::from_u32(self.function_manifest.len() as u32),
names: vec![export_name.to_string()],
});
self.function_manifest.push(FunctionSpec::new(
import_fn_ptr.as_usize() as u64,
0u32,
0u64,
0u64,
));
self
}
pub fn with_table_func(mut self, table_idx: u32, func_idx: u32, func: FunctionPointer) -> Self {
self.func_table.insert((table_idx, func_idx), func);
self
}
pub fn with_start_func(mut self, func: FunctionPointer) -> Self {
self.start_func = Some(func);
self
}
pub fn build(self) -> Arc<dyn Module> {
assert!(
self.sparse_page_data.len() * 4096 <= self.heap_spec.initial_size as usize,
"heap must fit in heap spec initial size"
);
let table_elements = self
.table_elements
.into_iter()
.enumerate()
.map(|(expected_idx, (idx, te))| {
assert_eq!(
idx, expected_idx,
"table element indices must be contiguous starting from 0"
);
te
})
.collect();
let globals_spec = self
.globals
.into_iter()
.enumerate()
.map(|(expected_idx, (idx, gs))| {
assert_eq!(
idx, expected_idx,
"global indices must be contiguous starting from 0"
);
gs
})
.collect();
let owned_module_data = OwnedModuleData::new(
Some(OwnedLinearMemorySpec {
heap: self.heap_spec,
initializer: OwnedSparseData::new(self.sparse_page_data)
.expect("sparse data pages are valid"),
}),
globals_spec,
self.function_info.clone(),
self.imports,
self.exports,
self.signatures,
ModuleFeatures::none(),
);
let serialized_module_data = owned_module_data
.to_ref()
.serialize()
.expect("serialization of module_data succeeds");
let module_data = ModuleData::deserialize(&serialized_module_data)
.map(|md| unsafe { std::mem::transmute(md) })
.expect("module data can be deserialized");
let mock = MockModule {
serialized_module_data,
module_data,
table_elements,
export_funcs: self.export_funcs,
func_table: self.func_table,
start_func: self.start_func,
function_manifest: self.function_manifest,
};
Arc::new(mock)
}
}
pub struct MockModule {
#[allow(dead_code)]
serialized_module_data: Vec<u8>,
module_data: ModuleData<'static>,
pub table_elements: Vec<TableElement>,
pub export_funcs: HashMap<&'static str, FunctionPointer>,
pub func_table: HashMap<(u32, u32), FunctionPointer>,
pub start_func: Option<FunctionPointer>,
pub function_manifest: Vec<FunctionSpec>,
}
unsafe impl Send for MockModule {}
unsafe impl Sync for MockModule {}
impl Module for MockModule {}
impl ModuleInternal for MockModule {
fn is_instruction_count_instrumented(&self) -> bool {
self.module_data.features().instruction_count
}
fn heap_spec(&self) -> Option<&HeapSpec> {
self.module_data.heap_spec()
}
fn globals(&self) -> &[GlobalSpec<'_>] {
self.module_data.globals_spec()
}
fn get_sparse_page_data(&self, page: usize) -> Option<&[u8]> {
if let Some(ref sparse_data) = self.module_data.sparse_data() {
*sparse_data.get_page(page)
} else {
None
}
}
fn sparse_page_data_len(&self) -> usize {
self.module_data.sparse_data().map(|d| d.len()).unwrap_or(0)
}
fn table_elements(&self) -> Result<&[TableElement], Error> {
Ok(&self.table_elements)
}
fn get_export_func(&self, sym: &str) -> Result<FunctionHandle, Error> {
let ptr = *self
.export_funcs
.get(sym)
.ok_or(Error::SymbolNotFound(sym.to_string()))?;
Ok(self.function_handle_from_ptr(ptr))
}
fn get_func_from_idx(&self, table_id: u32, func_id: u32) -> Result<FunctionHandle, Error> {
let ptr = self
.func_table
.get(&(table_id, func_id))
.cloned()
.ok_or(Error::FuncNotFound(table_id, func_id))?;
Ok(self.function_handle_from_ptr(ptr))
}
fn get_start_func(&self) -> Result<Option<FunctionHandle>, Error> {
Ok(self
.start_func
.map(|start| self.function_handle_from_ptr(start)))
}
fn function_manifest(&self) -> &[FunctionSpec] {
&self.function_manifest
}
fn addr_details(&self, _addr: *const c_void) -> Result<Option<AddrDetails>, Error> {
Ok(None)
}
fn get_signature(&self, fn_id: FunctionIndex) -> &Signature {
self.module_data.get_signature(fn_id)
}
}
pub struct MockExportBuilder {
sym: &'static str,
func: FunctionPointer,
func_len: Option<usize>,
traps: Option<&'static [TrapSite]>,
sig: Signature,
}
impl MockExportBuilder {
pub fn new(name: &'static str, func: FunctionPointer) -> MockExportBuilder {
MockExportBuilder {
sym: name,
func: func,
func_len: None,
traps: None,
sig: Signature {
params: vec![],
ret_ty: None,
},
}
}
pub fn with_func_len(mut self, len: usize) -> MockExportBuilder {
self.func_len = Some(len);
self
}
pub fn with_traps(mut self, traps: &'static [TrapSite]) -> MockExportBuilder {
self.traps = Some(traps);
self
}
pub fn with_sig(mut self, sig: Signature) -> MockExportBuilder {
self.sig = sig;
self
}
pub fn sym(&self) -> &'static str {
self.sym
}
pub fn func(&self) -> FunctionPointer {
self.func
}
pub fn func_len(&self) -> usize {
self.func_len.unwrap_or(1)
}
pub fn traps(&self) -> &'static [TrapSite] {
self.traps.unwrap_or(&[])
}
pub fn sig(&self) -> Signature {
self.sig.clone()
}
}