Merge pull request #2834 from rtfeldman/wasm-list-map-n

Wasm List.mapN
This commit is contained in:
Brian Carroll 2022-04-20 14:47:04 +01:00 committed by GitHub
commit f8156ffd53
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 181 additions and 84 deletions

View file

@ -3,6 +3,7 @@ use roc_builtins::bitcode::{self, FloatWidth, IntWidth};
use roc_error_macros::internal_error;
use roc_module::low_level::LowLevel;
use roc_module::symbol::Symbol;
use roc_mono::code_gen_help::HelperOp;
use roc_mono::ir::{HigherOrderLowLevel, PassedFunction, ProcLayout};
use roc_mono::layout::{Builtin, Layout, UnionLayout};
use roc_mono::low_level::HigherOrder;
@ -1014,58 +1015,73 @@ pub fn call_higher_order_lowlevel<'a>(
};
let wrapper_fn_idx = backend.register_helper_proc(wrapper_sym, wrapper_layout, source);
let inc_fn_idx = backend.gen_refcount_inc_for_zig(closure_data_layout);
let wrapper_fn_ptr = backend.get_fn_table_index(wrapper_fn_idx);
let inc_fn_ptr = backend.get_fn_table_index(inc_fn_idx);
let inc_fn_ptr = match closure_data_layout {
Layout::Struct {
field_layouts: &[], ..
} => {
// Our code gen would ignore the Unit arg, but the Zig builtin passes a pointer for it!
// That results in an exception (type signature mismatch in indirect call).
// The workaround is to use I32 layout, treating the (ignored) pointer as an integer.
backend.get_refcount_fn_ptr(Layout::Builtin(Builtin::Int(IntWidth::I32)), HelperOp::Inc)
}
_ => backend.get_refcount_fn_ptr(closure_data_layout, HelperOp::Inc),
};
match op {
// List.map : List elem_x, (elem_x -> elem_ret) -> List elem_ret
ListMap { xs } => {
let list_layout_in = backend.storage.symbol_layouts[xs];
ListMap { xs } => list_map_n(
bitcode::LIST_MAP,
backend,
&[*xs],
return_sym,
*return_layout,
wrapper_fn_ptr,
inc_fn_ptr,
closure_data_exists,
*captured_environment,
*owns_captured_environment,
),
let (elem_x, elem_ret) = match (list_layout_in, return_layout) {
(
Layout::Builtin(Builtin::List(elem_x)),
Layout::Builtin(Builtin::List(elem_ret)),
) => (elem_x, elem_ret),
_ => unreachable!("invalid layout for List.map arguments"),
};
let elem_x_size = elem_x.stack_size(TARGET_INFO);
let (elem_ret_size, elem_ret_align) = elem_ret.stack_size_and_alignment(TARGET_INFO);
ListMap2 { xs, ys } => list_map_n(
bitcode::LIST_MAP2,
backend,
&[*xs, *ys],
return_sym,
*return_layout,
wrapper_fn_ptr,
inc_fn_ptr,
closure_data_exists,
*captured_environment,
*owns_captured_environment,
),
let cb = &mut backend.code_builder;
ListMap3 { xs, ys, zs } => list_map_n(
bitcode::LIST_MAP3,
backend,
&[*xs, *ys, *zs],
return_sym,
*return_layout,
wrapper_fn_ptr,
inc_fn_ptr,
closure_data_exists,
*captured_environment,
*owns_captured_environment,
),
// Load return pointer & argument values
// Wasm signature: (i32, i64, i64, i32, i32, i32, i32, i32, i32, i32) -> nil
backend.storage.load_symbols(cb, &[return_sym]);
backend.storage.load_symbol_zig(cb, *xs); // list with capacity = 2 x i64 args
cb.i32_const(wrapper_fn_ptr);
if closure_data_exists {
backend.storage.load_symbols(cb, &[*captured_environment]);
} else {
// Normally, a zero-size arg would be eliminated in code gen, but Zig expects one!
cb.i32_const(0); // null pointer
}
cb.i32_const(inc_fn_ptr);
cb.i32_const(*owns_captured_environment as i32);
cb.i32_const(elem_ret_align as i32); // used for allocating the new list
cb.i32_const(elem_x_size as i32);
cb.i32_const(elem_ret_size as i32);
ListMap4 { xs, ys, zs, ws } => list_map_n(
bitcode::LIST_MAP4,
backend,
&[*xs, *ys, *zs, *ws],
return_sym,
*return_layout,
wrapper_fn_ptr,
inc_fn_ptr,
closure_data_exists,
*captured_environment,
*owns_captured_environment,
),
let num_wasm_args = 10; // 1 return pointer + 8 Zig args + list 2nd i64
let has_return_val = false;
backend.call_zig_builtin_after_loading_args(
bitcode::LIST_MAP,
num_wasm_args,
has_return_val,
);
}
ListMap2 { .. }
| ListMap3 { .. }
| ListMap4 { .. }
| ListMapWithIndex { .. }
ListMapWithIndex { .. }
| ListKeepIf { .. }
| ListWalk { .. }
| ListWalkUntil { .. }
@ -1079,3 +1095,71 @@ pub fn call_higher_order_lowlevel<'a>(
| DictWalk { .. } => todo!("{:?}", op),
}
}
fn unwrap_list_elem_layout(list_layout: Layout<'_>) -> &Layout<'_> {
match list_layout {
Layout::Builtin(Builtin::List(x)) => x,
e => internal_error!("expected List layout, got {:?}", e),
}
}
#[allow(clippy::too_many_arguments)]
fn list_map_n<'a>(
zig_fn_name: &'static str,
backend: &mut WasmBackend<'a>,
arg_symbols: &[Symbol],
return_sym: Symbol,
return_layout: Layout<'a>,
wrapper_fn_ptr: i32,
inc_fn_ptr: i32,
closure_data_exists: bool,
captured_environment: Symbol,
owns_captured_environment: bool,
) {
let arg_elem_layouts = Vec::from_iter_in(
arg_symbols
.iter()
.map(|sym| *unwrap_list_elem_layout(backend.storage.symbol_layouts[sym])),
backend.env.arena,
);
let elem_ret = unwrap_list_elem_layout(return_layout);
let (elem_ret_size, elem_ret_align) = elem_ret.stack_size_and_alignment(TARGET_INFO);
let cb = &mut backend.code_builder;
backend.storage.load_symbols(cb, &[return_sym]);
for s in arg_symbols {
backend.storage.load_symbol_zig(cb, *s);
}
cb.i32_const(wrapper_fn_ptr);
if closure_data_exists {
backend.storage.load_symbols(cb, &[captured_environment]);
} else {
// load_symbols assumes that a zero-size arg should be eliminated in code gen,
// but that's a specialization that our Zig code doesn't have! Pass a null pointer.
cb.i32_const(0);
}
cb.i32_const(inc_fn_ptr);
cb.i32_const(owns_captured_environment as i32);
cb.i32_const(elem_ret_align as i32);
for el in arg_elem_layouts.iter() {
cb.i32_const(el.stack_size(TARGET_INFO) as i32);
}
cb.i32_const(elem_ret_size as i32);
// If we have lists of different lengths, we may need to decrement
let num_wasm_args = if arg_elem_layouts.len() > 1 {
for el in arg_elem_layouts.iter() {
let ptr = backend.get_refcount_fn_ptr(*el, HelperOp::Dec);
backend.code_builder.i32_const(ptr);
}
7 + arg_elem_layouts.len() * 4
} else {
7 + arg_elem_layouts.len() * 3
};
let has_return_val = false;
backend.call_zig_builtin_after_loading_args(zig_fn_name, num_wasm_args, has_return_val);
}