wrap by-pointer functions so they own all arguments

cli/tests/cli_run.rs

@@ -227,7 +227,6 @@ mod cli_run {
         );
     }
 
-    #[ignore]
     #[test]
     #[serial(closure1)]
     fn closure1() {
@@ -240,7 +239,6 @@ mod cli_run {
         );
     }
 
-    #[ignore]
     #[test]
     #[serial(closure2)]
     fn closure2() {
@@ -253,8 +251,8 @@ mod cli_run {
         );
     }
 
-    #[ignore]
     #[test]
+    #[ignore]
     #[serial(closure3)]
     fn closure3() {
         check_output(
@@ -266,7 +264,6 @@ mod cli_run {
         );
     }
 
-    #[ignore]
     #[test]
     #[serial(closure4)]
     fn closure4() {

compiler/mono/src/borrow.rs

@@ -16,21 +16,12 @@ fn should_borrow_layout(layout: &Layout) -> bool {
 pub fn infer_borrow<'a>(
     arena: &'a Bump,
     procs: &MutMap<(Symbol, Layout<'a>), Proc<'a>>,
-    passed_by_pointer: &MutMap<(Symbol, Layout<'a>), Symbol>,
+    _passed_by_pointer: &MutMap<(Symbol, Layout<'a>), Symbol>,
 ) -> ParamMap<'a> {
     let mut param_map = ParamMap {
         items: MutMap::default(),
     };
 
-    for (key, other) in passed_by_pointer {
-        if let Some(proc) = procs.get(key) {
-            let mut proc: Proc = proc.clone();
-            proc.name = *other;
-
-            param_map.visit_proc_always_owned(arena, &proc, key.clone());
-        }
-    }
-
     for (key, proc) in procs {
         param_map.visit_proc(arena, proc, key.clone());
     }
@@ -159,6 +150,10 @@ impl<'a> ParamMap<'a> {
     }
 
     fn visit_proc(&mut self, arena: &'a Bump, proc: &Proc<'a>, key: (Symbol, Layout<'a>)) {
+        if proc.must_own_arguments {
+            self.visit_proc_always_owned(arena, proc, key);
+            return;
+        }
         let already_in_there = self.items.insert(
             Key::Declaration(proc.name, key.1),
             Self::init_borrow_args(arena, proc.args),

compiler/mono/src/ir.rs

@@ -118,6 +118,7 @@ pub struct Proc<'a> {
     pub closure_data_layout: Option<Layout<'a>>,
     pub ret_layout: Layout<'a>,
     pub is_self_recursive: SelfRecursive,
+    pub must_own_arguments: bool,
     pub host_exposed_layouts: HostExposedLayouts<'a>,
 }
 
@@ -1935,6 +1936,7 @@ fn specialize_external<'a>(
                 closure_data_layout,
                 ret_layout: full_layout,
                 is_self_recursive: recursivity,
+                must_own_arguments: false,
                 host_exposed_layouts,
             };
 
@@ -2071,6 +2073,7 @@ fn specialize_external<'a>(
                 closure_data_layout,
                 ret_layout,
                 is_self_recursive: recursivity,
+                must_own_arguments: false,
                 host_exposed_layouts,
             };
 
@@ -5707,11 +5710,26 @@ fn call_by_pointer<'a>(
     symbol: Symbol,
     layout: Layout<'a>,
 ) -> Expr<'a> {
+    // when we call a known function by-pointer, we must make sure we call a function that owns all
+    // its arguments (in an RC sense). we can't know this at this point, so we wrap such calls in
+    // a proc that we guarantee owns all its arguments. E.g. we turn
+    //
+    // foo = \x -> ...
+    //
+    // x = List.map [ ... ] foo
+    //
+    // into
+    //
+    // foo = \x -> ...
+    //
+    // @owns_all_arguments
+    // foo1 = \x -> foo x
+    //
+    // x = List.map [ ... ] foo1
     if env.is_imported_symbol(symbol) || procs.partial_procs.contains_key(&symbol) {
-        // println!("going to specialize {:?} {:?}", symbol, layout);
-
         match layout {
             Layout::FunctionPointer(arg_layouts, ret_layout) => {
+                if arg_layouts.iter().any(|l| l.contains_refcounted()) {
                     let name = env.unique_symbol();
                     let mut args = Vec::with_capacity_in(arg_layouts.len(), env.arena);
                     let mut arg_symbols = Vec::with_capacity_in(arg_layouts.len(), env.arena);
@@ -5724,9 +5742,8 @@ fn call_by_pointer<'a>(
                     let args = args.into_bump_slice();
 
                     let call_symbol = env.unique_symbol();
-                let callable_symbol = env.unique_symbol();
-                let call_type = CallType::ByPointer {
-                    name: callable_symbol,
+                    let call_type = CallType::ByName {
+                        name: symbol,
                         full_layout: layout.clone(),
                         ret_layout: ret_layout.clone(),
                         arg_layouts,
@@ -5741,14 +5758,6 @@ fn call_by_pointer<'a>(
 
                     body = Stmt::Let(call_symbol, expr, ret_layout.clone(), env.arena.alloc(body));
 
-                let fpointer = Expr::FunctionPointer(symbol, layout.clone());
-                body = Stmt::Let(
-                    callable_symbol,
-                    fpointer,
-                    ret_layout.clone(),
-                    env.arena.alloc(body),
-                );
-
                     let closure_data_layout = None;
                     let proc = Proc {
                         name,
@@ -5757,6 +5766,7 @@ fn call_by_pointer<'a>(
                         closure_data_layout,
                         ret_layout: ret_layout.clone(),
                         is_self_recursive: SelfRecursive::NotSelfRecursive,
+                        must_own_arguments: true,
                         host_exposed_layouts: HostExposedLayouts::NotHostExposed,
                     };
 
@@ -5764,6 +5774,11 @@ fn call_by_pointer<'a>(
                         .specialized
                         .insert((name, layout.clone()), InProgressProc::Done(proc));
                     Expr::FunctionPointer(name, layout)
+                } else {
+                    // if none of the arguments is refcounted, then owning the arguments has no
+                    // meaning
+                    Expr::FunctionPointer(symbol, layout)
+                }
             }
             _ => unreachable!("not a function pointer layout"),
         }

examples/benchmarks/Closure2.roc

@@ -12,4 +12,6 @@ main =
 
     Task.succeed {}
         |> Task.map (\_ -> x)
-        |> Task.map (\_ -> {})
+        |> Task.map toUnit
+
+toUnit = (\_ -> {})