hook borrow inference up

crates/compiler/mono/src/inc_dec.rs

@@ -15,7 +15,7 @@ use roc_module::low_level::LowLevel;
 use roc_module::{low_level::LowLevelWrapperType, symbol::Symbol};
 
 use crate::ir::ErasedField;
-use crate::layout::LambdaName;
+use crate::layout::{LambdaName, Niche};
 use crate::{
     ir::{
         BranchInfo, Call, CallType, Expr, HigherOrderLowLevel, JoinPointId, ListLiteralElement,
@@ -34,15 +34,14 @@ pub fn insert_inc_dec_operations<'a>(
     procedures: &mut HashMap<(Symbol, ProcLayout<'a>), Proc<'a>, BuildHasherDefault<WyHash>>,
 ) {
     // TODO remove this clone?
-    let x = procedures.clone();
-    let ps = arena.alloc(x);
+    let ps = arena.alloc(procedures.clone());
 
     let borrow_signatures = crate::borrow::infer_borrow_signatures(arena, layout_interner, ps);
     let borrow_signatures = arena.alloc(borrow_signatures);
 
-    for ((s, _), sig) in borrow_signatures.iter() {
-        dbg!((s, sig));
-    }
+    //    for ((s, _), sig) in borrow_signatures.iter() {
+    //        dbg!((s, sig));
+    //    }
 
     // All calls to lowlevels are wrapped in another function to help with type inference and return/parameter layouts.
     // But this lowlevel might get inlined into the caller of the wrapper and thus removing any reference counting operations.
@@ -257,8 +256,6 @@ Type containing data about the symbols consumption of a join point.
 */
 type JoinPointConsumption = MutSet<Symbol>;
 
-type Key<'a> = (LambdaName<'a>, InLayout<'a>, &'a [InLayout<'a>]);
-
 /**
 The environment for the reference counting pass.
 Contains the symbols rc types and the ownership.
@@ -322,9 +319,13 @@ impl<'v> RefcountEnvironment<'v> {
     Add a symbol to the environment if it is reference counted.
     */
     fn add_symbol(&mut self, symbol: Symbol) {
+        self.add_symbol_with(symbol, Ownership::Owned)
+    }
+
+    fn add_symbol_with(&mut self, symbol: Symbol, ownership: Ownership) {
         match self.get_symbol_rc_type(&symbol) {
             VarRcType::ReferenceCounted => {
-                self.symbols_ownership.insert(symbol, Ownership::Owned);
+                self.symbols_ownership.insert(symbol, ownership);
             }
             VarRcType::NotReferenceCounted => {
                 // If this symbol is not reference counted, we don't need to do anything.
@@ -434,16 +435,21 @@ fn insert_inc_dec_operations_proc<'a>(
     };
 
     // Add all arguments to the environment (if they are reference counted)
-    let proc_symbols = proc.args.iter().map(|(_layout, symbol)| symbol);
-    for symbol in proc_symbols.clone() {
-        environment.add_symbol(*symbol);
+    let borrow_signature = borrow_signatures
+        .get(&(proc.name.name(), proc.proc_layout(arena)))
+        .unwrap();
+    for ((_, symbol), ownership) in proc.args.iter().zip(borrow_signature.iter()) {
+        environment.add_symbol_with(*symbol, ownership);
     }
 
     // Update the body with reference count statements.
     let new_body = insert_refcount_operations_stmt(arena, &mut environment, &proc.body);
 
     // Insert decrement statements for unused parameters (which are still marked as owned).
-    let rc_proc_symbols = proc_symbols
+    let rc_proc_symbols = proc
+        .args
+        .iter()
+        .map(|(_layout, symbol)| symbol)
         .filter(|symbol| environment.symbols_ownership.contains_key(symbol))
         .copied()
         .collect_in::<Vec<_>>(arena);
@@ -738,10 +744,11 @@ fn insert_refcount_operations_stmt<'v, 'a>(
         } => {
             // Assuming that the values in the closure of the body of this jointpoint are already bound.
             // Assuming that all symbols are still owned. (So that we can determine what symbols got consumed in the join point.)
-            debug_assert!(environment
-                .symbols_ownership
-                .iter()
-                .all(|(_, ownership)| ownership.is_owned()));
+
+            //            debug_assert!(environment
+            //                .symbols_ownership
+            //                .iter()
+            //                .all(|(_, ownership)| ownership.is_owned()));
 
             let mut body_env = environment.clone();
 
@@ -996,12 +1003,45 @@ fn insert_refcount_operations_binding<'a>(
             call_type,
         }) => {
             match call_type.clone().replace_lowlevel_wrapper() {
-                // A by name call refers to a normal function call.
-                // Normal functions take all their parameters as owned, so we can mark them all as such.
-                CallType::ByName { .. } => {
-                    let new_let = new_let!(stmt);
+                CallType::ByName {
+                    name,
+                    arg_layouts,
+                    ret_layout,
+                    ..
+                } => {
+                    //                    let new_let = new_let!(stmt);
+                    //
+                    //                    inc_owned!(arguments.iter().copied(), new_let)
 
-                    inc_owned!(arguments.iter().copied(), new_let)
+                    let proc_layout = ProcLayout {
+                        arguments: arg_layouts,
+                        result: ret_layout,
+                        niche: Niche::NONE,
+                    };
+
+                    let borrow_signature = match environment
+                        .borrow_signatures
+                        .get(&(name.name(), proc_layout))
+                    {
+                        Some(s) => s,
+                        None => unreachable!("no borrow signature for {name:?} layout"),
+                    };
+
+                    let owned_arguments = arguments
+                        .iter()
+                        .copied()
+                        .zip(borrow_signature.iter())
+                        .filter_map(|(symbol, ownership)| ownership.is_owned().then_some(symbol));
+                    let borrowed_arguments = arguments
+                        .iter()
+                        .copied()
+                        .zip(borrow_signature.iter())
+                        .filter_map(|(symbol, ownership)| {
+                            ownership.is_borrowed().then_some(symbol)
+                        });
+                    let new_stmt = dec_borrowed!(borrowed_arguments, stmt);
+                    let new_let = new_let!(new_stmt);
+                    inc_owned!(owned_arguments, new_let)
                 }
                 // A normal Roc function call, but we don't actually know where its target is.
                 // As such, we assume that it takes all parameters as owned, as will the function