roc/compiler/test_gen/src/helpers/wasm.rs
2021-11-30 10:58:03 +00:00

420 lines
12 KiB
Rust

use bumpalo::collections::vec::Vec;
use bumpalo::Bump;
use std::cell::Cell;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use tempfile::{tempdir, TempDir};
use crate::helpers::from_wasm32_memory::FromWasm32Memory;
use crate::helpers::wasm32_test_result::Wasm32TestResult;
use roc_builtins::bitcode;
use roc_can::builtins::builtin_defs_map;
use roc_collections::all::{MutMap, MutSet};
use roc_gen_wasm::wasm_module::linking::{WasmObjectSymbol, WASM_SYM_UNDEFINED};
use roc_gen_wasm::wasm_module::sections::{Import, ImportDesc};
use roc_gen_wasm::wasm_module::{
CodeBuilder, Export, ExportType, LocalId, Signature, SymInfo, ValueType, WasmModule,
};
use roc_gen_wasm::MEMORY_NAME;
#[allow(unused_imports)]
use roc_mono::ir::PRETTY_PRINT_IR_SYMBOLS;
const TEST_WRAPPER_NAME: &str = "test_wrapper";
std::thread_local! {
static TEST_COUNTER: Cell<u32> = Cell::new(0);
}
fn promote_expr_to_module(src: &str) -> String {
let mut buffer = String::from("app \"test\" provides [ main ] to \"./platform\"\n\nmain =\n");
for line in src.lines() {
// indent the body!
buffer.push_str(" ");
buffer.push_str(line);
buffer.push('\n');
}
buffer
}
#[allow(dead_code)]
pub fn helper_wasm<'a, T: Wasm32TestResult>(
arena: &'a bumpalo::Bump,
src: &str,
stdlib: &'a roc_builtins::std::StdLib,
_result_type_dummy: &T,
) -> wasmer::Instance {
use std::path::{Path, PathBuf};
let filename = PathBuf::from("Test.roc");
let src_dir = Path::new("fake/test/path");
let module_src;
let temp;
if src.starts_with("app") {
// this is already a module
module_src = src;
} else {
// this is an expression, promote it to a module
temp = promote_expr_to_module(src);
module_src = &temp;
}
let exposed_types = MutMap::default();
let ptr_bytes = 4;
let loaded = roc_load::file::load_and_monomorphize_from_str(
arena,
filename,
module_src,
stdlib,
src_dir,
exposed_types,
ptr_bytes,
builtin_defs_map,
);
let loaded = loaded.expect("failed to load module");
use roc_load::file::MonomorphizedModule;
let MonomorphizedModule {
module_id,
procedures,
mut interns,
exposed_to_host,
..
} = loaded;
// You can comment and uncomment this block out to get more useful information
// while you're working on the wasm backend!
{
// println!("=========== Procedures ==========");
// if PRETTY_PRINT_IR_SYMBOLS {
// println!("");
// for proc in procedures.values() {
// println!("{}", proc.to_pretty(200));
// }
// } else {
// println!("{:?}", procedures.values());
// }
// println!("=================================\n");
// println!("=========== Interns ==========");
// println!("{:?}", interns);
// println!("=================================\n");
// println!("=========== Exposed ==========");
// println!("{:?}", exposed_to_host);
// println!("=================================\n");
}
debug_assert_eq!(exposed_to_host.len(), 1);
let exposed_to_host = exposed_to_host.keys().copied().collect::<MutSet<_>>();
let env = roc_gen_wasm::Env {
arena,
module_id,
exposed_to_host,
};
let (mut wasm_module, main_fn_index) =
roc_gen_wasm::build_module_help(&env, &mut interns, procedures).unwrap();
T::insert_test_wrapper(arena, &mut wasm_module, TEST_WRAPPER_NAME, main_fn_index);
// We can either generate the test platform or write an external source file, whatever works
generate_test_platform(&mut wasm_module, arena);
let mut module_bytes = std::vec::Vec::with_capacity(4096);
wasm_module.serialize_mut(&mut module_bytes);
// now, do wasmer stuff
use wasmer::{Instance, Module, Store};
let store = Store::default();
// Keep the final .wasm file for debugging with wasm-objdump or wasm2wat
const DEBUG_WASM_FILE: bool = false;
let wasmer_module = {
let tmp_dir: TempDir; // directory for normal test runs, deleted when dropped
let debug_dir: String; // persistent directory for debugging
let wasm_build_dir: &Path = if DEBUG_WASM_FILE {
// Directory name based on a hash of the Roc source
let mut hash_state = DefaultHasher::new();
src.hash(&mut hash_state);
let src_hash = hash_state.finish();
debug_dir = format!("/tmp/roc/gen_wasm/{:016x}", src_hash);
std::fs::create_dir_all(&debug_dir).unwrap();
println!(
"Debug command:\n\twasm-objdump -sdx {}/final.wasm",
&debug_dir
);
Path::new(&debug_dir)
} else {
tmp_dir = tempdir().unwrap();
tmp_dir.path()
};
let final_wasm_file = wasm_build_dir.join("final.wasm");
let app_o_file = wasm_build_dir.join("app.o");
let libc_a_file = "../gen_wasm/lib/libc.a";
// write the module to a file so the linker can access it
std::fs::write(&app_o_file, &module_bytes).unwrap();
let args = &[
"wasm-ld",
// input files
app_o_file.to_str().unwrap(),
bitcode::BUILTINS_WASM32_OBJ_PATH,
libc_a_file,
// output
"-o",
final_wasm_file.to_str().unwrap(),
// we don't define `_start`
"--no-entry",
// If you only specify test_wrapper, it will stop at the call to UserApp_main_1
// But if you specify both exports, you get all the dependencies.
//
// It seems that it will not write out an export you didn't explicitly specify,
// even if it's a dependency of another export!
// In our case we always export main and test_wrapper so that's OK.
"--export",
"test_wrapper",
"--export",
"#UserApp_main_1",
];
let linker_output = std::process::Command::new("zig")
.args(args)
.output()
.unwrap();
if !linker_output.status.success() {
print!("\nLINKER FAILED\n");
for arg in args {
print!("{} ", arg);
}
println!("\n{}", std::str::from_utf8(&linker_output.stdout).unwrap());
println!("{}", std::str::from_utf8(&linker_output.stderr).unwrap());
}
Module::from_file(&store, &final_wasm_file).unwrap()
};
// First, we create the `WasiEnv`
use wasmer_wasi::WasiState;
let mut wasi_env = WasiState::new("hello").finalize().unwrap();
// Then, we get the import object related to our WASI
// and attach it to the Wasm instance.
let import_object = wasi_env
.import_object(&wasmer_module)
.unwrap_or_else(|_| wasmer::imports!());
Instance::new(&wasmer_module, &import_object).unwrap()
}
#[allow(dead_code)]
pub fn assert_wasm_evals_to_help<T>(src: &str, expected: T) -> Result<T, String>
where
T: FromWasm32Memory + Wasm32TestResult,
{
let arena = bumpalo::Bump::new();
// NOTE the stdlib must be in the arena; just taking a reference will segfault
let stdlib = arena.alloc(roc_builtins::std::standard_stdlib());
let instance = crate::helpers::wasm::helper_wasm(&arena, src, stdlib, &expected);
let memory = instance.exports.get_memory(MEMORY_NAME).unwrap();
let test_wrapper = instance.exports.get_function(TEST_WRAPPER_NAME).unwrap();
match test_wrapper.call(&[]) {
Err(e) => Err(format!("{:?}", e)),
Ok(result) => {
let address = match result[0] {
wasmer::Value::I32(a) => a,
_ => panic!(),
};
let output = <T as FromWasm32Memory>::decode(memory, address as u32);
Ok(output)
}
}
}
#[allow(unused_macros)]
macro_rules! assert_wasm_evals_to {
($src:expr, $expected:expr, $ty:ty, $transform:expr) => {
match $crate::helpers::wasm::assert_wasm_evals_to_help::<$ty>($src, $expected) {
Err(msg) => panic!("{:?}", msg),
Ok(actual) => {
assert_eq!($transform(actual), $expected)
}
}
};
($src:expr, $expected:expr, $ty:ty) => {
$crate::helpers::wasm::assert_wasm_evals_to!(
$src,
$expected,
$ty,
$crate::helpers::wasm::identity
);
};
($src:expr, $expected:expr, $ty:ty, $transform:expr) => {
$crate::helpers::wasm::assert_wasm_evals_to!($src, $expected, $ty, $transform);
};
}
#[allow(unused_macros)]
macro_rules! assert_evals_to {
($src:expr, $expected:expr, $ty:ty) => {{
assert_evals_to!($src, $expected, $ty, $crate::helpers::wasm::identity);
}};
($src:expr, $expected:expr, $ty:ty, $transform:expr) => {
// Same as above, except with an additional transformation argument.
{
$crate::helpers::wasm::assert_wasm_evals_to!($src, $expected, $ty, $transform);
}
};
}
#[allow(dead_code)]
pub fn identity<T>(value: T) -> T {
value
}
#[allow(unused_imports)]
pub(crate) use assert_evals_to;
#[allow(unused_imports)]
pub(crate) use assert_wasm_evals_to;
fn wrap_libc_fn<'a>(
module: &mut WasmModule<'a>,
arena: &'a Bump,
roc_name: &'a str,
libc_name: &'a str,
params: &'a [(ValueType, bool)],
ret_type: Option<ValueType>,
) {
let symbol_table = module.linking.symbol_table_mut();
// Type signatures
let mut wrapper_signature = Signature {
param_types: Vec::with_capacity_in(params.len(), arena),
ret_type,
};
let mut libc_signature = Signature {
param_types: Vec::with_capacity_in(params.len(), arena),
ret_type,
};
for (ty, used) in params.iter() {
wrapper_signature.param_types.push(*ty);
if *used {
libc_signature.param_types.push(*ty);
}
}
/*
* Import a function from libc
*/
let libc_signature_index = module.types.insert(libc_signature);
// Import
let import_index = module.import.entries.len() as u32;
module.import.entries.push(Import {
module: "env",
name: libc_name.to_string(),
description: ImportDesc::Func {
signature_index: libc_signature_index,
},
});
// Linker info
let libc_sym_idx = symbol_table.len() as u32;
symbol_table.push(SymInfo::Function(WasmObjectSymbol::Imported {
flags: WASM_SYM_UNDEFINED,
index: import_index,
}));
/*
* Export a wrapper function
*/
// Declaration
let wrapper_sig_index = module.types.insert(wrapper_signature);
module.function.signature_indices.push(wrapper_sig_index);
// Body
let mut code_builder = CodeBuilder::new(arena);
let mut num_libc_args = 0;
for (i, (_, used)) in params.iter().enumerate() {
if *used {
code_builder.get_local(LocalId(i as u32));
num_libc_args += 1;
}
}
code_builder.call(
import_index,
libc_sym_idx,
num_libc_args,
ret_type.is_some(),
);
code_builder.build_fn_header(&[], 0, None);
let wrapper_index = module.code.code_builders.len() as u32;
module.code.code_builders.push(code_builder);
// Export
module.export.entries.push(Export {
name: roc_name.to_string(),
ty: ExportType::Func,
index: wrapper_index,
});
// Linker symbol
symbol_table.push(SymInfo::Function(WasmObjectSymbol::Defined {
flags: 0,
index: wrapper_index,
name: roc_name.to_string(),
}));
}
fn generate_test_platform<'a>(module: &mut WasmModule<'a>, arena: &'a Bump) {
use ValueType::I32;
wrap_libc_fn(
module,
arena,
"roc_alloc",
"malloc",
// only the first argument of roc_alloc is passed to malloc
&[(I32, true), (I32, false)],
Some(I32),
);
wrap_libc_fn(
module,
arena,
"roc_dealloc",
"free",
&[(I32, true), (I32, false)],
None,
);
wrap_libc_fn(
module,
arena,
"roc_realloc",
"realloc",
&[(I32, true), (I32, false), (I32, true), (I32, false)],
Some(I32),
);
}