WIP

compiler/mono/src/alias_analysis.rs

@@ -605,6 +605,201 @@ fn call_spec(
             *update_mode,
             call.arguments,
         ),
+        NewHigherOrderLowLevel {
+            specialization_id,
+            closure_env_layout,
+            op,
+            arg_layouts,
+            ret_layout,
+            function_name,
+            function_env,
+            ..
+        } => {
+            let array = specialization_id.to_bytes();
+            let spec_var = CalleeSpecVar(&array);
+
+            let it = arg_layouts.iter().copied();
+            let bytes = func_name_bytes_help(*function_name, it, *ret_layout);
+            let name = FuncName(&bytes);
+            let module = MOD_APP;
+
+            use crate::low_level::HigherOrder::*;
+
+            match op {
+                DictWalk { xs, state } => {
+                    let dict = env.symbols[xs];
+                    let state = env.symbols[state];
+                    let closure_env = env.symbols[function_env];
+
+                    let bag = builder.add_get_tuple_field(block, dict, DICT_BAG_INDEX)?;
+                    let _cell = builder.add_get_tuple_field(block, dict, DICT_CELL_INDEX)?;
+
+                    let first = builder.add_bag_get(block, bag)?;
+
+                    let key = builder.add_get_tuple_field(block, first, 0)?;
+                    let val = builder.add_get_tuple_field(block, first, 1)?;
+
+                    let argument = if closure_env_layout.is_none() {
+                        builder.add_make_tuple(block, &[state, key, val])?
+                    } else {
+                        builder.add_make_tuple(block, &[state, key, val, closure_env])?
+                    };
+                    builder.add_call(block, spec_var, module, name, argument)?;
+                }
+
+                ListWalk { xs, state }
+                | ListWalkBackwards { xs, state }
+                | ListWalkUntil { xs, state } => {
+                    let list = env.symbols[xs];
+                    let state = env.symbols[state];
+                    let closure_env = env.symbols[function_env];
+
+                    let bag = builder.add_get_tuple_field(block, list, LIST_BAG_INDEX)?;
+                    let _cell = builder.add_get_tuple_field(block, list, LIST_CELL_INDEX)?;
+
+                    let first = builder.add_bag_get(block, bag)?;
+
+                    let argument = if closure_env_layout.is_none() {
+                        builder.add_make_tuple(block, &[state, first])?
+                    } else {
+                        builder.add_make_tuple(block, &[state, first, closure_env])?
+                    };
+                    builder.add_call(block, spec_var, module, name, argument)?;
+                }
+
+                ListMapWithIndex { xs } => {
+                    let list = env.symbols[xs];
+                    let closure_env = env.symbols[function_env];
+
+                    let bag = builder.add_get_tuple_field(block, list, LIST_BAG_INDEX)?;
+                    let _cell = builder.add_get_tuple_field(block, list, LIST_CELL_INDEX)?;
+
+                    let first = builder.add_bag_get(block, bag)?;
+                    let index = builder.add_make_tuple(block, &[])?;
+
+                    let argument = if closure_env_layout.is_none() {
+                        builder.add_make_tuple(block, &[index, first])?
+                    } else {
+                        builder.add_make_tuple(block, &[index, first, closure_env])?
+                    };
+                    builder.add_call(block, spec_var, module, name, argument)?;
+                }
+
+                ListMap { xs } => {
+                    let list = env.symbols[xs];
+                    let closure_env = env.symbols[function_env];
+
+                    let bag1 = builder.add_get_tuple_field(block, list, LIST_BAG_INDEX)?;
+                    let _cell1 = builder.add_get_tuple_field(block, list, LIST_CELL_INDEX)?;
+
+                    let elem1 = builder.add_bag_get(block, bag1)?;
+
+                    let argument = if closure_env_layout.is_none() {
+                        builder.add_make_tuple(block, &[elem1])?
+                    } else {
+                        builder.add_make_tuple(block, &[elem1, closure_env])?
+                    };
+                    builder.add_call(block, spec_var, module, name, argument)?;
+                }
+
+                ListSortWith { xs } => {
+                    let list = env.symbols[xs];
+                    let closure_env = env.symbols[function_env];
+
+                    let bag1 = builder.add_get_tuple_field(block, list, LIST_BAG_INDEX)?;
+                    let _cell1 = builder.add_get_tuple_field(block, list, LIST_CELL_INDEX)?;
+
+                    let elem1 = builder.add_bag_get(block, bag1)?;
+
+                    let argument = if closure_env_layout.is_none() {
+                        builder.add_make_tuple(block, &[elem1, elem1])?
+                    } else {
+                        builder.add_make_tuple(block, &[elem1, elem1, closure_env])?
+                    };
+                    builder.add_call(block, spec_var, module, name, argument)?;
+                }
+
+                ListMap2 { xs, ys } => {
+                    let list1 = env.symbols[xs];
+                    let list2 = env.symbols[ys];
+                    let closure_env = env.symbols[function_env];
+
+                    let bag1 = builder.add_get_tuple_field(block, list1, LIST_BAG_INDEX)?;
+                    let _cell1 = builder.add_get_tuple_field(block, list1, LIST_CELL_INDEX)?;
+                    let elem1 = builder.add_bag_get(block, bag1)?;
+
+                    let bag2 = builder.add_get_tuple_field(block, list2, LIST_BAG_INDEX)?;
+                    let _cell2 = builder.add_get_tuple_field(block, list2, LIST_CELL_INDEX)?;
+                    let elem2 = builder.add_bag_get(block, bag2)?;
+
+                    let argument = if closure_env_layout.is_none() {
+                        builder.add_make_tuple(block, &[elem1, elem2])?
+                    } else {
+                        builder.add_make_tuple(block, &[elem1, elem2, closure_env])?
+                    };
+                    builder.add_call(block, spec_var, module, name, argument)?;
+                }
+
+                ListMap3 { xs, ys, zs } => {
+                    let list1 = env.symbols[xs];
+                    let list2 = env.symbols[ys];
+                    let list3 = env.symbols[zs];
+                    let closure_env = env.symbols[function_env];
+
+                    let bag1 = builder.add_get_tuple_field(block, list1, LIST_BAG_INDEX)?;
+                    let _cell1 = builder.add_get_tuple_field(block, list1, LIST_CELL_INDEX)?;
+                    let elem1 = builder.add_bag_get(block, bag1)?;
+
+                    let bag2 = builder.add_get_tuple_field(block, list2, LIST_BAG_INDEX)?;
+                    let _cell2 = builder.add_get_tuple_field(block, list2, LIST_CELL_INDEX)?;
+                    let elem2 = builder.add_bag_get(block, bag2)?;
+
+                    let bag3 = builder.add_get_tuple_field(block, list3, LIST_BAG_INDEX)?;
+                    let _cell3 = builder.add_get_tuple_field(block, list3, LIST_CELL_INDEX)?;
+                    let elem3 = builder.add_bag_get(block, bag3)?;
+
+                    let argument = if closure_env_layout.is_none() {
+                        builder.add_make_tuple(block, &[elem1, elem2, elem3])?
+                    } else {
+                        builder.add_make_tuple(block, &[elem1, elem2, elem3, closure_env])?
+                    };
+                    builder.add_call(block, spec_var, module, name, argument)?;
+                }
+
+                ListKeepIf { xs } | ListKeepOks { xs } | ListKeepErrs { xs } => {
+                    let list = env.symbols[xs];
+                    let closure_env = env.symbols[function_env];
+
+                    let bag = builder.add_get_tuple_field(block, list, LIST_BAG_INDEX)?;
+                    // let _cell = builder.add_get_tuple_field(block, list, LIST_CELL_INDEX)?;
+
+                    let first = builder.add_bag_get(block, bag)?;
+
+                    let argument = if closure_env_layout.is_none() {
+                        builder.add_make_tuple(block, &[first])?
+                    } else {
+                        builder.add_make_tuple(block, &[first, closure_env])?
+                    };
+                    let result = builder.add_call(block, spec_var, module, name, argument)?;
+                    let unit = builder.add_tuple_type(&[])?;
+                    builder.add_unknown_with(block, &[result], unit)?;
+                }
+            }
+
+            // TODO overly pessimstic
+            // filter_map because one of the arguments is a function name, which
+            // is not defined in the env
+            let arguments: Vec<_> = call
+                .arguments
+                .iter()
+                .filter_map(|symbol| env.symbols.get(symbol))
+                .copied()
+                .collect();
+
+            let result_type = layout_spec(builder, layout)?;
+
+            builder.add_unknown_with(block, &arguments, result_type)
+        }
         HigherOrderLowLevel {
             specialization_id,
             closure_env_layout,

compiler/mono/src/borrow.rs

@@ -593,6 +593,91 @@ impl<'a> BorrowInfState<'a> {
                 self.own_args_using_bools(arguments, ps);
             }
 
+            NewHigherOrderLowLevel {
+                op,
+                arg_layouts,
+                ret_layout,
+                function_name,
+                function_env,
+                ..
+            } => {
+                use crate::low_level::HigherOrder::*;
+
+                let closure_layout = ProcLayout {
+                    arguments: arg_layouts,
+                    result: *ret_layout,
+                };
+
+                let function_ps = match param_map.get_symbol(*function_name, closure_layout) {
+                    Some(function_ps) => function_ps,
+                    None => unreachable!(),
+                };
+
+                match op {
+                    ListMap { xs }
+                    | ListKeepIf { xs }
+                    | ListKeepOks { xs }
+                    | ListKeepErrs { xs } => {
+                        // own the list if the function wants to own the element
+                        if !function_ps[0].borrow {
+                            self.own_var(*xs);
+                        }
+                    }
+                    ListMapWithIndex { xs } => {
+                        // own the list if the function wants to own the element
+                        if !function_ps[1].borrow {
+                            self.own_var(*xs);
+                        }
+                    }
+                    ListMap2 { xs, ys } => {
+                        // own the lists if the function wants to own the element
+                        if !function_ps[0].borrow {
+                            self.own_var(*xs);
+                        }
+
+                        if !function_ps[1].borrow {
+                            self.own_var(*ys);
+                        }
+                    }
+                    ListMap3 { xs, ys, zs } => {
+                        // own the lists if the function wants to own the element
+                        if !function_ps[0].borrow {
+                            self.own_var(*xs);
+                        }
+                        if !function_ps[1].borrow {
+                            self.own_var(*ys);
+                        }
+                        if !function_ps[2].borrow {
+                            self.own_var(*zs);
+                        }
+                    }
+                    ListSortWith { xs } => {
+                        // always own the input list
+                        self.own_var(*xs);
+                    }
+                    ListWalk { xs, state }
+                    | ListWalkUntil { xs, state }
+                    | ListWalkBackwards { xs, state }
+                    | DictWalk { xs, state } => {
+                        // own the default value if the function wants to own it
+                        if !function_ps[0].borrow {
+                            self.own_var(*state);
+                        }
+
+                        // own the data structure if the function wants to own the element
+                        if !function_ps[1].borrow {
+                            self.own_var(*xs);
+                        }
+                    }
+                }
+
+                // own the closure environment if the function needs to own it
+                let function_env_position = op.function_arity();
+                if let Some(false) = function_ps.get(function_env_position).map(|p| p.borrow) {
+                    self.own_var(*function_env);
+                }
+            }
+
             HigherOrderLowLevel {
                 op,
                 arg_layouts,
@@ -952,7 +1037,7 @@ pub fn lowlevel_borrow_signature(arena: &Bump, op: LowLevel) -> &[bool] {
     use LowLevel::*;
 
     // TODO is true or false more efficient for non-refcounted layouts?
-    let irrelevant = OWNED;
+    let irrelevant = BORROWED;
     let function = irrelevant;
     let closure_data = irrelevant;
     let owned = OWNED;
@@ -1071,6 +1156,7 @@ fn call_info_call<'a>(call: &crate::ir::Call<'a>, info: &mut CallInfo<'a>) {
         Foreign { .. } => {}
         LowLevel { .. } => {}
         HigherOrderLowLevel { .. } => {}
+        NewHigherOrderLowLevel { .. } => {}
     }
 }
 

compiler/mono/src/inc_dec.rs

@@ -463,6 +463,163 @@ impl<'a> Context<'a> {
                 &*self.arena.alloc(Stmt::Let(z, v, l, b))
             }
 
+            NewHigherOrderLowLevel {
+                op,
+                closure_env_layout,
+                specialization_id,
+                arg_layouts,
+                ret_layout,
+                function_name,
+                function_env,
+                ..
+            } => {
+                // setup
+                use crate::low_level::HigherOrder::*;
+
+                macro_rules! create_call {
+                    ($borrows:expr) => {
+                        Expr::Call(crate::ir::Call {
+                            call_type: if let Some(OWNED) = $borrows.map(|p| p.borrow) {
+                                NewHigherOrderLowLevel {
+                                    op: *op,
+                                    closure_env_layout: *closure_env_layout,
+                                    function_owns_closure_data: true,
+                                    specialization_id: *specialization_id,
+                                    function_name: *function_name,
+                                    function_env: *function_env,
+                                    arg_layouts,
+                                    ret_layout: *ret_layout,
+                                }
+                            } else {
+                                call_type
+                            },
+                            arguments,
+                        })
+                    };
+                }
+
+                macro_rules! decref_if_owned {
+                    ($borrows:expr, $argument:expr, $stmt:expr) => {
+                        if !$borrows {
+                            self.arena.alloc(Stmt::Refcounting(
+                                ModifyRc::DecRef($argument),
+                                self.arena.alloc($stmt),
+                            ))
+                        } else {
+                            $stmt
+                        }
+                    };
+                }
+
+                const FUNCTION: bool = BORROWED;
+                const CLOSURE_DATA: bool = BORROWED;
+
+                let function_layout = ProcLayout {
+                    arguments: arg_layouts,
+                    result: *ret_layout,
+                };
+
+                let function_ps = match self.param_map.get_symbol(*function_name, function_layout) {
+                    Some(function_ps) => function_ps,
+                    None => unreachable!(),
+                };
+
+                match op {
+                    ListMap { xs }
+                    | ListKeepIf { xs }
+                    | ListKeepOks { xs }
+                    | ListKeepErrs { xs } => {
+                        let borrows = [function_ps[0].borrow, FUNCTION, CLOSURE_DATA];
+
+                        let b = self.add_dec_after_lowlevel(arguments, &borrows, b, b_live_vars);
+
+                        // if the list is owned, then all elements have been consumed, but not the list itself
+                        let b = decref_if_owned!(function_ps[0].borrow, *xs, b);
+
+                        let v = create_call!(function_ps.get(1));
+
+                        &*self.arena.alloc(Stmt::Let(z, v, l, b))
+                    }
+                    ListMap2 { xs, ys } => {
+                        let borrows = [
+                            function_ps[0].borrow,
+                            function_ps[1].borrow,
+                            FUNCTION,
+                            CLOSURE_DATA,
+                        ];
+
+                        let b = self.add_dec_after_lowlevel(arguments, &borrows, b, b_live_vars);
+
+                        let b = decref_if_owned!(function_ps[0].borrow, *xs, b);
+                        let b = decref_if_owned!(function_ps[1].borrow, *ys, b);
+
+                        let v = create_call!(function_ps.get(2));
+
+                        &*self.arena.alloc(Stmt::Let(z, v, l, b))
+                    }
+                    ListMap3 { xs, ys, zs } => {
+                        let borrows = [
+                            function_ps[0].borrow,
+                            function_ps[1].borrow,
+                            function_ps[2].borrow,
+                            FUNCTION,
+                            CLOSURE_DATA,
+                        ];
+
+                        let b = self.add_dec_after_lowlevel(arguments, &borrows, b, b_live_vars);
+
+                        let b = decref_if_owned!(function_ps[0].borrow, *xs, b);
+                        let b = decref_if_owned!(function_ps[1].borrow, *ys, b);
+                        let b = decref_if_owned!(function_ps[2].borrow, *zs, b);
+
+                        let v = create_call!(function_ps.get(3));
+
+                        &*self.arena.alloc(Stmt::Let(z, v, l, b))
+                    }
+                    ListMapWithIndex { xs } => {
+                        let borrows = [function_ps[1].borrow, FUNCTION, CLOSURE_DATA];
+
+                        let b = self.add_dec_after_lowlevel(arguments, &borrows, b, b_live_vars);
+
+                        let b = decref_if_owned!(function_ps[1].borrow, *xs, b);
+
+                        let v = create_call!(function_ps.get(2));
+
+                        &*self.arena.alloc(Stmt::Let(z, v, l, b))
+                    }
+                    ListSortWith { xs: _ } => {
+                        let borrows = [OWNED, FUNCTION, CLOSURE_DATA];
+
+                        let b = self.add_dec_after_lowlevel(arguments, &borrows, b, b_live_vars);
+
+                        let v = create_call!(function_ps.get(2));
+
+                        &*self.arena.alloc(Stmt::Let(z, v, l, b))
+                    }
+                    ListWalk { xs, state: _ }
+                    | ListWalkUntil { xs, state: _ }
+                    | ListWalkBackwards { xs, state: _ }
+                    | DictWalk { xs, state: _ } => {
+                        // borrow data structure based on first argument of the folded function
+                        // borrow the default based on second argument of the folded function
+                        let borrows = [
+                            function_ps[1].borrow,
+                            function_ps[0].borrow,
+                            FUNCTION,
+                            CLOSURE_DATA,
+                        ];
+
+                        let b = self.add_dec_after_lowlevel(arguments, &borrows, b, b_live_vars);
+
+                        let b = decref_if_owned!(function_ps[1].borrow, *xs, b);
+
+                        let v = create_call!(function_ps.get(2));
+
+                        &*self.arena.alloc(Stmt::Let(z, v, l, b))
+                    }
+                }
+            }
+
             HigherOrderLowLevel {
                 op,
                 closure_env_layout,

compiler/mono/src/ir.rs

@@ -1050,6 +1050,13 @@ impl<'a> Call<'a> {
                     .text(format!("lowlevel {:?} ", lowlevel))
                     .append(alloc.intersperse(it, " "))
             }
+            NewHigherOrderLowLevel { op: lowlevel, .. } => {
+                let it = arguments.iter().map(|s| symbol_to_doc(alloc, *s));
+
+                alloc
+                    .text(format!("lowlevel {:?} ", lowlevel))
+                    .append(alloc.intersperse(it, " "))
+            }
             Foreign {
                 ref foreign_symbol, ..
             } => {
@@ -1113,6 +1120,27 @@ pub enum CallType<'a> {
         arg_layouts: &'a [Layout<'a>],
         ret_layout: Layout<'a>,
     },
+    NewHigherOrderLowLevel {
+        op: crate::low_level::HigherOrder,
+        /// the layout of the closure argument, if any
+        closure_env_layout: Option<Layout<'a>>,
+
+        /// name of the top-level function that is passed as an argument
+        /// e.g. in `List.map xs Num.abs` this would be `Num.abs`
+        function_name: Symbol,
+
+        /// Symbol of the environment captured by the function argument
+        function_env: Symbol,
+
+        /// does the function argument need to own the closure data
+        function_owns_closure_data: bool,
+
+        /// specialization id of the function argument, used for name generation
+        specialization_id: CallSpecId,
+        /// function layout, used for name generation
+        arg_layouts: &'a [Layout<'a>],
+        ret_layout: Layout<'a>,
+    },
 }
 
 #[derive(Clone, Debug, PartialEq)]
@@ -5447,6 +5475,7 @@ fn substitute_in_call<'a>(
         CallType::Foreign { .. } => None,
         CallType::LowLevel { .. } => None,
         CallType::HigherOrderLowLevel { .. } => None,
+        CallType::NewHigherOrderLowLevel { .. } => None,
     };
 
     let mut did_change = false;

compiler/mono/src/lib.rs

@@ -8,6 +8,7 @@ pub mod expand_rc;
 pub mod inc_dec;
 pub mod ir;
 pub mod layout;
+pub mod low_level;
 pub mod reset_reuse;
 pub mod tail_recursion;