move List.map* into roc

This commit is contained in:
Brendan Hansknecht 2024-07-12 15:32:49 -07:00
parent 7d8fbfbe85
commit c734a27b59
No known key found for this signature in database
GPG key ID: 0EA784685083E75B
16 changed files with 60 additions and 1445 deletions

View file

@ -39,7 +39,6 @@ pub enum ProcSource {
Roc,
Helper,
/// Wrapper function for higher-order calls from Zig to Roc
HigherOrderMapper(usize),
HigherOrderCompare(usize),
}
@ -491,126 +490,6 @@ impl<'a, 'r> WasmBackend<'a, 'r> {
self.module.names.append_function(wasm_fn_index, name);
}
/// Build a wrapper around a Roc procedure so that it can be called from Zig builtins List.map*
///
/// The generic Zig code passes *pointers* to all of the argument values (e.g. on the heap in a List).
/// Numbers up to 64 bits are passed by value, so we need to load them from the provided pointer.
/// Everything else is passed by reference, so we can just pass the pointer through.
///
/// NOTE: If the builtins expected the return pointer first and closure data last, we could eliminate the wrapper
/// when all args are pass-by-reference and non-zero size. But currently we need it to swap those around.
pub fn build_higher_order_mapper(
&mut self,
wrapper_lookup_idx: usize,
inner_lookup_idx: usize,
) {
use Align::*;
use ValueType::*;
let ProcLookupData {
name: wrapper_name,
layout: wrapper_proc_layout,
..
} = self.proc_lookup[wrapper_lookup_idx];
let wrapper_arg_layouts = wrapper_proc_layout.arguments;
// Our convention is that the last arg of the wrapper is the heap return pointer
let heap_return_ptr_id = LocalId(wrapper_arg_layouts.len() as u32 - 1);
let inner_ret_layout = match wrapper_arg_layouts
.last()
.map(|l| self.layout_interner.get_repr(*l))
{
Some(LayoutRepr::Ptr(inner)) => WasmLayout::new(self.layout_interner, inner),
x => internal_error!("Higher-order wrapper: invalid return layout {:?}", x),
};
let ret_type_and_size = match inner_ret_layout.return_method() {
ReturnMethod::NoReturnValue => None,
ReturnMethod::Primitive(ty, size) => {
// If the inner function returns a primitive, load the address to store it at
// After the call, it will be under the call result in the value stack
self.code_builder.get_local(heap_return_ptr_id);
Some((ty, size))
}
ReturnMethod::WriteToPointerArg => {
// If the inner function writes to a return pointer, load its address
self.code_builder.get_local(heap_return_ptr_id);
None
}
};
// Load all the arguments for the inner function
for (i, wrapper_arg) in wrapper_arg_layouts.iter().enumerate() {
let is_closure_data = i == 0; // Skip closure data (first for wrapper, last for inner). We'll handle it below.
let is_return_pointer = i == wrapper_arg_layouts.len() - 1; // Skip return pointer (may not be an arg for inner. And if it is, swaps from end to start)
if is_closure_data || is_return_pointer {
continue;
}
let inner_layout = match self.layout_interner.get_repr(*wrapper_arg) {
LayoutRepr::Ptr(inner) => inner,
x => internal_error!("Expected a Ptr layout, got {:?}", x),
};
if self.layout_interner.stack_size(inner_layout) == 0 {
continue;
}
// Load the argument pointer. If it's a primitive value, dereference it too.
self.code_builder.get_local(LocalId(i as u32));
self.dereference_boxed_value(inner_layout);
}
// If the inner function has closure data, it's the last arg of the inner fn
let closure_data_layout = wrapper_arg_layouts[0];
if self.layout_interner.stack_size(closure_data_layout) > 0 {
// The closure data exists, and will have been passed in to the wrapper as a
// one-element struct.
let inner_closure_data_layout = match self.layout_interner.get_repr(closure_data_layout)
{
LayoutRepr::Struct([inner]) => inner,
other => internal_error!(
"Expected a boxed layout for wrapped closure data, got {:?}",
other
),
};
self.code_builder.get_local(LocalId(0));
// Since the closure data is wrapped in a one-element struct, we've been passed in the
// pointer to that struct in the stack memory. To get the closure data we just need to
// dereference the pointer.
self.dereference_boxed_value(*inner_closure_data_layout);
}
// Call the wrapped inner function
let inner_wasm_fn_index = self.fn_index_offset + inner_lookup_idx as u32;
self.code_builder.call(inner_wasm_fn_index);
// If the inner function returns a primitive, store it to the address we loaded at the very beginning
if let Some((ty, size)) = ret_type_and_size {
match (ty, size) {
(I64, 8) => self.code_builder.i64_store(Bytes8, 0),
(I32, 4) => self.code_builder.i32_store(Bytes4, 0),
(I32, 2) => self.code_builder.i32_store16(Bytes2, 0),
(I32, 1) => self.code_builder.i32_store8(Bytes1, 0),
(F32, 4) => self.code_builder.f32_store(Bytes4, 0),
(F64, 8) => self.code_builder.f64_store(Bytes8, 0),
_ => {
internal_error!("Cannot store {:?} with alignment of {:?}", ty, size);
}
}
}
// Write empty function header (local variables array with zero length)
self.code_builder.build_fn_header_and_footer(&[], 0, None);
self.module.add_function_signature(Signature {
param_types: bumpalo::vec![in self.env.arena; I32; wrapper_arg_layouts.len()],
ret_type: None,
});
self.append_proc_debug_name(wrapper_name);
self.reset();
}
/// Build a wrapper around a Roc comparison proc so that it can be called from higher-order Zig builtins.
/// Comparison procedure signature is: closure_data, a, b -> Order (u8)
///

View file

@ -178,7 +178,6 @@ pub fn build_app_module<'a, 'r>(
match source {
Roc => { /* already generated */ }
Helper => backend.build_proc(helper_iter.next().unwrap()),
HigherOrderMapper(inner_idx) => backend.build_higher_order_mapper(idx, *inner_idx),
HigherOrderCompare(inner_idx) => backend.build_higher_order_compare(idx, *inner_idx),
}
}

View file

@ -377,7 +377,7 @@ impl<'a> LowLevelCall<'a> {
backend.call_host_fn_after_loading_args(bitcode::LIST_DECREF);
}
ListMap | ListMap2 | ListMap3 | ListMap4 | ListSortWith => {
ListSortWith => {
internal_error!("HigherOrder lowlevels should not be handled here")
}
@ -2614,7 +2614,7 @@ fn num_is_finite(backend: &mut WasmBackend<'_, '_>, argument: Symbol) {
pub fn call_higher_order_lowlevel<'a>(
backend: &mut WasmBackend<'a, '_>,
return_sym: Symbol,
return_layout: &InLayout<'a>,
_return_layout: &InLayout<'a>,
higher_order: &'a HigherOrderLowLevel<'a>,
) {
use HigherOrder::*;
@ -2726,9 +2726,6 @@ pub fn call_higher_order_lowlevel<'a>(
.unwrap();
match op {
ListSortWith { .. } => ProcSource::HigherOrderCompare(passed_proc_index),
ListMap { .. } | ListMap2 { .. } | ListMap3 { .. } | ListMap4 { .. } => {
ProcSource::HigherOrderMapper(passed_proc_index)
}
}
};
let wrapper_sym = backend.create_symbol(&format!("#wrap#{fn_name:?}"));
@ -2753,19 +2750,6 @@ pub fn call_higher_order_lowlevel<'a>(
wrapper_arg_layouts.extend(boxed_closure_arg_layouts);
match helper_proc_source {
ProcSource::HigherOrderMapper(_) => {
// Our convention for mappers is that they write to the heap via the last argument
wrapper_arg_layouts.push(
backend
.layout_interner
.insert_direct_no_semantic(LayoutRepr::Ptr(*result_layout)),
);
ProcLayout {
arguments: wrapper_arg_layouts.into_bump_slice(),
result: Layout::UNIT,
niche: fn_name.niche(),
}
}
ProcSource::HigherOrderCompare(_) => ProcLayout {
arguments: wrapper_arg_layouts.into_bump_slice(),
result: *result_layout,
@ -2792,58 +2776,6 @@ pub fn call_higher_order_lowlevel<'a>(
};
match op {
ListMap { xs } => list_map_n(
bitcode::LIST_MAP,
backend,
&[*xs],
return_sym,
*return_layout,
wrapper_fn_ptr,
inc_n_fn_ptr,
closure_data_exists,
wrapped_captured_environment,
*owns_captured_environment,
),
ListMap2 { xs, ys } => list_map_n(
bitcode::LIST_MAP2,
backend,
&[*xs, *ys],
return_sym,
*return_layout,
wrapper_fn_ptr,
inc_n_fn_ptr,
closure_data_exists,
wrapped_captured_environment,
*owns_captured_environment,
),
ListMap3 { xs, ys, zs } => list_map_n(
bitcode::LIST_MAP3,
backend,
&[*xs, *ys, *zs],
return_sym,
*return_layout,
wrapper_fn_ptr,
inc_n_fn_ptr,
closure_data_exists,
wrapped_captured_environment,
*owns_captured_environment,
),
ListMap4 { xs, ys, zs, ws } => list_map_n(
bitcode::LIST_MAP4,
backend,
&[*xs, *ys, *zs, *ws],
return_sym,
*return_layout,
wrapper_fn_ptr,
inc_n_fn_ptr,
closure_data_exists,
wrapped_captured_environment,
*owns_captured_environment,
),
ListSortWith { xs } => {
let elem_in_layout = unwrap_list_elem_layout(
backend
@ -2905,63 +2837,6 @@ fn unwrap_list_elem_layout(list_layout: LayoutRepr) -> InLayout {
}
}
#[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: InLayout<'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
.layout_interner
.get_repr(backend.storage.symbol_layouts[sym]),
)
}),
backend.env.arena,
);
let elem_in_ret = unwrap_list_elem_layout(backend.layout_interner.get_repr(return_layout));
let elem_ret = backend.layout_interner.get_repr(elem_in_ret);
let (elem_ret_size, elem_ret_align) =
elem_ret.stack_size_and_alignment(backend.layout_interner);
let elem_ret_refcounted = backend.layout_interner.contains_refcounted(elem_in_ret);
let cb = &mut backend.code_builder;
let mut args_vec = Vec::with_capacity_in(arg_symbols.len() + 1, backend.env.arena);
args_vec.push(return_sym);
args_vec.extend_from_slice(arg_symbols);
backend.storage.load_symbols(cb, &args_vec);
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(backend.layout_interner.stack_size(*el) as i32);
}
cb.i32_const(elem_ret_size as i32);
backend.code_builder.i32_const(elem_ret_refcounted as i32);
backend.call_host_fn_after_loading_args(zig_fn_name);
}
fn ensure_symbol_is_in_memory<'a>(
backend: &mut WasmBackend<'a, '_>,
symbol: Symbol,