hook borrow inference up

crates/compiler/mono/src/borrow.rs

@@ -16,8 +16,6 @@ impl std::fmt::Debug for BorrowSignature {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let mut f = &mut f.debug_struct("BorrowSignature");
 
-        dbg!(self.0);
-
         for (i, ownership) in self.iter().enumerate() {
             f = f.field(&format!("_{i}"), &ownership);
         }
@@ -59,7 +57,7 @@ impl BorrowSignature {
         }
     }
 
-    fn iter(&self) -> impl Iterator<Item = Ownership> + '_ {
+    pub fn iter(&self) -> impl Iterator<Item = Ownership> + '_ {
         let mut i = 0;
 
         std::iter::from_fn(move || {
@@ -80,23 +78,25 @@ impl std::ops::Index<usize> for BorrowSignature {
 
 pub(crate) type BorrowSignatures<'a> = MutMap<(Symbol, ProcLayout<'a>), BorrowSignature>;
 
-pub(crate) fn infer_borrow_signatures<'a>(
+pub(crate) fn infer_borrow_signatures<'a, 'b: 'a>(
     arena: &'a Bump,
     interner: &impl LayoutInterner<'a>,
-    procs: &'a MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>,
+    procs: &'b MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>,
 ) -> BorrowSignatures<'a> {
-    let host_exposed_procs = &[];
+    // let host_exposed_procs = &[];
 
-    let mut borrow_signatures = procs
+    let mut borrow_signatures: BorrowSignatures = procs
         .iter()
-        .map(|(key, proc)| {
+        .map(|(_key, proc)| {
             let mut signature = BorrowSignature::new(proc.args.len());
 
+            let key = (proc.name.name(), proc.proc_layout(arena));
+
             for (i, in_layout) in key.1.arguments.iter().enumerate() {
                 signature.set(i, layout_to_ownership(*in_layout, interner));
             }
 
-            (*key, signature)
+            (key, signature)
         })
         .collect();
 
@@ -107,8 +107,6 @@ pub(crate) fn infer_borrow_signatures<'a>(
     let matrix = construct_reference_matrix(arena, procs);
     let sccs = matrix.strongly_connected_components_all();
 
-    let mut env = ();
-
     for (group, _) in sccs.groups() {
         // This is a fixed-point analysis
         //
@@ -117,21 +115,49 @@ pub(crate) fn infer_borrow_signatures<'a>(
         // when that doesn't lead to conflicts the change is kept, otherwise it may be reverted
         //
         // when the signatures no longer change, the analysis stops and returns the signatures
+
+        //        // initialize borrow signatures for everyone
+        //        for index in group.iter_ones() {
+        //            let (key, proc) = procs.iter().nth(index).unwrap();
+        //
+        //            if proc.args.is_empty() {
+        //                continue;
+        //            }
+        //
+        //            // host-exposed functions must always own their arguments.
+        //            let is_host_exposed = host_exposed_procs.contains(&key.0);
+        //
+        //            let key = (proc.name.name(), proc.proc_layout(arena));
+        //
+        //            // initialize the borrow signature based on the layout if first time
+        //            borrow_signatures.entry(key).or_insert_with(|| {
+        //                let mut borrow_signature = BorrowSignature::new(proc.args.len());
+        //
+        //                for (i, in_layout) in key.1.arguments.iter().enumerate() {
+        //                    borrow_signature.set(i, layout_to_ownership(*in_layout, interner));
+        //                }
+        //
+        //                borrow_signature
+        //            });
+        //        }
+
         loop {
             for index in group.iter_ones() {
-                let (key, proc) = procs.iter().nth(index).unwrap();
+                let (_, proc) = procs.iter().nth(index).unwrap();
+                let key = (proc.name.name(), proc.proc_layout(arena));
 
                 if proc.args.is_empty() {
                     continue;
                 }
 
-                // host-exposed functions must always own their arguments.
+                let mut state = State {
-                let is_host_exposed = host_exposed_procs.contains(&key.0);
+                    args: proc.args,
+                    borrow_signature: *borrow_signatures.get(&key).unwrap(),
+                };
 
-                let mut state = State::new(arena, interner, &mut borrow_signatures, proc);
                 state.inspect_stmt(&mut borrow_signatures, &proc.body);
 
-                borrow_signatures.insert(*key, state.borrow_signature);
+                borrow_signatures.insert(key, state.borrow_signature);
             }
 
             // if there were no modifications, we're done
@@ -234,13 +260,13 @@ impl<'a> State<'a> {
                 self.mark_owned(*s);
             }
             Stmt::Refcounting(_, _) => unreachable!("not inserted yet"),
-            Stmt::Expect { .. } | Stmt::ExpectFx { .. } => {
+            Stmt::Expect { remainder, .. } | Stmt::ExpectFx { remainder, .. } => {
-                // TODO do we rely on values being passed by-value here?
+                // based on my reading of inc_dec.rs, expect borrows the symbols
-                // it would be better to pass by-reference in general
+                self.inspect_stmt(borrow_signatures, remainder);
             }
-            Stmt::Dbg { .. } => {
+            Stmt::Dbg { remainder, .. } => {
-                // TODO do we rely on values being passed by-value here?
+                // based on my reading of inc_dec.rs, expect borrows the symbol
-                // it would be better to pass by-reference in general
+                self.inspect_stmt(borrow_signatures, remainder);
             }
             Stmt::Join {
                 body, remainder, ..
@@ -280,7 +306,7 @@ impl<'a> State<'a> {
 
                 let borrow_signature = match borrow_signatures.get(&(name.name(), proc_layout)) {
                     Some(s) => s,
-                    None => todo!("no borrow signature for function/layout"),
+                    None => unreachable!("no borrow signature for {name:?} layout"),
                 };
 
                 for (argument, ownership) in arguments.iter().zip(borrow_signature.iter()) {

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(procedures.clone());
-    let ps = arena.alloc(x);
 
     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() {
+    //    for ((s, _), sig) in borrow_signatures.iter() {
-        dbg!((s, sig));
+    //        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);
+    let borrow_signature = borrow_signatures
-    for symbol in proc_symbols.clone() {
+        .get(&(proc.name.name(), proc.proc_layout(arena)))
-        environment.add_symbol(*symbol);
+        .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
+            //            debug_assert!(environment
-                .iter()
+            //                .symbols_ownership
-                .all(|(_, ownership)| ownership.is_owned()));
+            //                .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.
+                CallType::ByName {
-                // Normal functions take all their parameters as owned, so we can mark them all as such.
+                    name,
-                CallType::ByName { .. } => {
+                    arg_layouts,
-                    let new_let = new_let!(stmt);
+                    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

crates/compiler/mono/src/ir.rs

@@ -398,7 +398,7 @@ impl<'a> Proc<'a> {
     }
 
     pub fn proc_layout(&self, arena: &'a Bump) -> ProcLayout<'a> {
-        let args = Vec::from_iter_in(self.args.iter().map(|(a, b)| *a), arena);
+        let args = Vec::from_iter_in(self.args.iter().map(|(a, _)| *a), arena);
 
         ProcLayout {
             arguments: args.into_bump_slice(),