Merge branch 'trunk' into cli-test-tweaks

Cargo.lock

@@ -1,5 +1,7 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
+version = 3
+
 [[package]]
 name = "ab_glyph"
 version = "0.2.11"
@@ -3356,6 +3358,7 @@ dependencies = [
  "roc_types",
  "roc_unify",
  "ven_ena",
+ "ven_graph",
  "ven_pretty",
 ]
 

compiler/build/src/program.rs

@@ -116,9 +116,10 @@ pub fn gen_from_mono_module(
         }
 
         if name.starts_with("roc_builtins.dict")
-            || name.starts_with("dict.RocDict")
             || name.starts_with("roc_builtins.list")
+            || name.starts_with("roc_builtins.dec")
             || name.starts_with("list.RocList")
+            || name.starts_with("dict.RocDict")
         {
             function.add_attribute(AttributeLoc::Function, enum_attr);
         }

compiler/module/src/symbol.rs

@@ -774,6 +774,9 @@ define_builtins! {
 
         // a caller (wrapper) for comparison
         21 GENERIC_COMPARE_REF: "#generic_compare_ref"
+
+        // used to initialize parameters in borrow.rs
+        22 EMPTY_PARAM: "#empty_param"
     }
     1 NUM: "Num" => {
         0 NUM_NUM: "Num" imported // the Num.Num type alias

compiler/mono/Cargo.toml

@@ -19,6 +19,7 @@ morphic_lib = { path = "../../vendor/morphic_lib" }
 bumpalo = { version = "3.6.1", features = ["collections"] }
 hashbrown = { version = "0.11.2", features = [ "bumpalo" ] }
 ven_ena = { path = "../../vendor/ena" }
+ven_graph = { path = "../../vendor/pathfinding" }
 linked-hash-map = "0.5.4"
 
 [dev-dependencies]

compiler/mono/src/alias_analysis.rs

@@ -842,7 +842,19 @@ fn lowlevel_spec(
 
             builder.add_bag_insert(block, bag, to_insert)?;
 
-            Ok(list)
+            let new_cell = builder.add_new_heap_cell(block)?;
+            builder.add_make_tuple(block, &[new_cell, bag])
+        }
+        ListSwap => {
+            let list = env.symbols[&arguments[0]];
+
+            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 _unit = builder.add_update(block, update_mode_var, cell)?;
+
+            let new_cell = builder.add_new_heap_cell(block)?;
+            builder.add_make_tuple(block, &[new_cell, bag])
         }
         ListAppend => {
             let list = env.symbols[&arguments[0]];
@@ -853,9 +865,11 @@ fn lowlevel_spec(
 
             let _unit = builder.add_update(block, update_mode_var, cell)?;
 
+            // TODO new heap cell
             builder.add_bag_insert(block, bag, to_insert)?;
 
-            Ok(list)
+            let new_cell = builder.add_new_heap_cell(block)?;
+            builder.add_make_tuple(block, &[new_cell, bag])
         }
         DictEmpty => {
             match layout {
@@ -887,7 +901,6 @@ fn lowlevel_spec(
             let cell = builder.add_get_tuple_field(block, dict, DICT_CELL_INDEX)?;
 
             let _unit = builder.add_touch(block, cell)?;
-
             builder.add_bag_get(block, bag)
         }
         DictInsert => {
@@ -904,7 +917,8 @@ fn lowlevel_spec(
 
             builder.add_bag_insert(block, bag, key_value)?;
 
-            Ok(dict)
+            let new_cell = builder.add_new_heap_cell(block)?;
+            builder.add_make_tuple(block, &[new_cell, bag])
         }
         _other => {
             // println!("missing {:?}", _other);

compiler/mono/src/borrow.rs

@@ -2,7 +2,7 @@ use crate::ir::{Expr, JoinPointId, Param, Proc, ProcLayout, Stmt};
 use crate::layout::Layout;
 use bumpalo::collections::Vec;
 use bumpalo::Bump;
-use roc_collections::all::{MutMap, MutSet};
+use roc_collections::all::{default_hasher, MutMap, MutSet};
 use roc_module::low_level::LowLevel;
 use roc_module::symbol::Symbol;
 
@@ -20,8 +20,22 @@ pub fn infer_borrow<'a>(
     arena: &'a Bump,
     procs: &MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>,
 ) -> ParamMap<'a> {
-    let mut param_map = ParamMap {
+    // intern the layouts
-        items: MutMap::default(),
+    let mut declaration_to_index = MutMap::with_capacity_and_hasher(procs.len(), default_hasher());
+
+    let mut param_map = {
+        let mut i = 0;
+        for key in procs.keys() {
+            declaration_to_index.insert(*key, ParamOffset(i));
+
+            i += key.1.arguments.len();
+        }
+
+        ParamMap {
+            declaration_to_index,
+            join_points: MutMap::default(),
+            declarations: bumpalo::vec![in arena; Param::EMPTY; i],
+        }
     };
 
     for (key, proc) in procs {
@@ -33,10 +47,94 @@ pub fn infer_borrow<'a>(
         param_set: MutSet::default(),
         owned: MutMap::default(),
         modified: false,
-        param_map,
         arena,
     };
 
+    // next we first partition the functions into strongly connected components, then do a
+    // topological sort on these components, finally run the fix-point borrow analysis on each
+    // component (in top-sorted order, from primitives (std-lib) to main)
+
+    let successor_map = &make_successor_mapping(arena, procs);
+    let successors = move |key: &Symbol| successor_map[key].iter().copied();
+
+    let mut symbols = Vec::with_capacity_in(procs.len(), arena);
+    symbols.extend(procs.keys().map(|x| x.0));
+
+    let sccs = ven_graph::strongly_connected_components(&symbols, successors);
+
+    let mut symbol_to_component = MutMap::default();
+    for (i, symbols) in sccs.iter().enumerate() {
+        for symbol in symbols {
+            symbol_to_component.insert(*symbol, i);
+        }
+    }
+
+    let mut component_to_successors = Vec::with_capacity_in(sccs.len(), arena);
+    for (i, symbols) in sccs.iter().enumerate() {
+        // guess: every function has ~1 successor
+        let mut succs = Vec::with_capacity_in(symbols.len(), arena);
+
+        for symbol in symbols {
+            for s in successors(symbol) {
+                let c = symbol_to_component[&s];
+
+                // don't insert self to prevent cycles
+                if c != i {
+                    succs.push(c);
+                }
+            }
+        }
+
+        succs.sort_unstable();
+        succs.dedup();
+
+        component_to_successors.push(succs);
+    }
+
+    let mut components = Vec::with_capacity_in(component_to_successors.len(), arena);
+    components.extend(0..component_to_successors.len());
+
+    let mut groups = Vec::new_in(arena);
+
+    let component_to_successors = &component_to_successors;
+    match ven_graph::topological_sort_into_groups(&components, |c: &usize| {
+        component_to_successors[*c].iter().copied()
+    }) {
+        Ok(component_groups) => {
+            let mut component_to_group = bumpalo::vec![in arena; usize::MAX; components.len()];
+
+            // for each component, store which group it is in
+            for (group_index, component_group) in component_groups.iter().enumerate() {
+                for component in component_group {
+                    component_to_group[*component] = group_index;
+                }
+            }
+
+            // prepare groups
+            groups.reserve(component_groups.len());
+            for _ in 0..component_groups.len() {
+                groups.push(Vec::new_in(arena));
+            }
+
+            for (key, proc) in procs {
+                let symbol = key.0;
+                let offset = param_map.declaration_to_index[key];
+
+                // the component this symbol is a part of
+                let component = symbol_to_component[&symbol];
+
+                // now find the group that this component belongs to
+                let group = component_to_group[component];
+
+                groups[group].push((proc, offset));
+            }
+        }
+        Err((_groups, _remainder)) => {
+            unreachable!("because we find strongly-connected components first");
+        }
+    }
+
+    for group in groups.into_iter().rev() {
         // This is a fixed-point analysis
         //
         // all functions initiall own all their parameters
@@ -45,12 +143,8 @@ pub fn infer_borrow<'a>(
         //
         // when the signatures no longer change, the analysis stops and returns the signatures
         loop {
-        // sort the symbols (roughly) in definition order.
+            for (proc, param_offset) in group.iter() {
-        // TODO in the future I think we need to do this properly, and group
+                env.collect_proc(&mut param_map, proc, *param_offset);
-        // mutually recursive functions (or just make all their arguments owned)
-
-        for (key, proc) in procs {
-            env.collect_proc(proc, key.1);
             }
 
             if !env.modified {
@@ -61,54 +155,50 @@ pub fn infer_borrow<'a>(
                 env.modified = false;
             }
         }
-
-    env.param_map
     }
 
-#[derive(Debug, PartialEq, Eq, Hash, Clone)]
+    param_map
-pub enum Key<'a> {
-    Declaration(Symbol, ProcLayout<'a>),
-    JoinPoint(JoinPointId),
 }
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
+pub struct ParamOffset(usize);
+
+impl From<ParamOffset> for usize {
+    fn from(id: ParamOffset) -> Self {
+        id.0 as usize
+    }
+}
+
+#[derive(Debug, Clone)]
 pub struct ParamMap<'a> {
-    items: MutMap<Key<'a>, &'a [Param<'a>]>,
+    /// Map a (Symbol, ProcLayout) pair to the starting index in the `declarations` array
-}
+    declaration_to_index: MutMap<(Symbol, ProcLayout<'a>), ParamOffset>,
-
+    /// the parameters of all functions in a single flat array.
-impl<'a> IntoIterator for ParamMap<'a> {
+    ///
-    type Item = (Key<'a>, &'a [Param<'a>]);
+    /// - the map above gives the index of the first parameter for the function
-    type IntoIter = <std::collections::HashMap<Key<'a>, &'a [Param<'a>]> as IntoIterator>::IntoIter;
+    /// - the length of the ProcLayout's argument field gives the total number of parameters
-
+    ///
-    fn into_iter(self) -> Self::IntoIter {
+    /// These can be read by taking a slice into this array, and can also be updated in-place
-        self.items.into_iter()
+    declarations: Vec<'a, Param<'a>>,
-    }
+    join_points: MutMap<JoinPointId, &'a [Param<'a>]>,
-}
-
-impl<'a> IntoIterator for &'a ParamMap<'a> {
-    type Item = (&'a Key<'a>, &'a &'a [Param<'a>]);
-    type IntoIter =
-        <&'a std::collections::HashMap<Key<'a>, &'a [Param<'a>]> as IntoIterator>::IntoIter;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.items.iter()
-    }
 }
 
 impl<'a> ParamMap<'a> {
-    pub fn get_symbol(&self, symbol: Symbol, layout: ProcLayout<'a>) -> Option<&'a [Param<'a>]> {
+    pub fn get_symbol(&self, symbol: Symbol, layout: ProcLayout<'a>) -> Option<&[Param<'a>]> {
-        let key = Key::Declaration(symbol, layout);
+        let index: usize = self.declaration_to_index[&(symbol, layout)].into();
 
-        self.items.get(&key).copied()
+        self.declarations.get(index..index + layout.arguments.len())
     }
     pub fn get_join_point(&self, id: JoinPointId) -> &'a [Param<'a>] {
-        let key = Key::JoinPoint(id);
+        match self.join_points.get(&id) {
-
-        match self.items.get(&key) {
             Some(slice) => slice,
             None => unreachable!("join point not in param map: {:?}", id),
         }
     }
+
+    pub fn iter_symbols(&'a self) -> impl Iterator<Item = &'a Symbol> {
+        self.declaration_to_index.iter().map(|t| &t.0 .0)
+    }
 }
 
 impl<'a> ParamMap<'a> {
@@ -156,11 +246,16 @@ impl<'a> ParamMap<'a> {
             self.visit_proc_always_owned(arena, proc, key);
             return;
         }
-        let already_in_there = self.items.insert(
+
-            Key::Declaration(proc.name, key.1),
+        let index: usize = self.declaration_to_index[&key].into();
-            Self::init_borrow_args(arena, proc.args),
+
-        );
+        for (i, param) in Self::init_borrow_args(arena, proc.args)
-        debug_assert!(already_in_there.is_none());
+            .iter()
+            .copied()
+            .enumerate()
+        {
+            self.declarations[index + i] = param;
+        }
 
         self.visit_stmt(arena, proc.name, &proc.body);
     }
@@ -171,11 +266,15 @@ impl<'a> ParamMap<'a> {
         proc: &Proc<'a>,
         key: (Symbol, ProcLayout<'a>),
     ) {
-        let already_in_there = self.items.insert(
+        let index: usize = self.declaration_to_index[&key].into();
-            Key::Declaration(proc.name, key.1),
+
-            Self::init_borrow_args_always_owned(arena, proc.args),
+        for (i, param) in Self::init_borrow_args_always_owned(arena, proc.args)
-        );
+            .iter()
-        debug_assert!(already_in_there.is_none());
+            .copied()
+            .enumerate()
+        {
+            self.declarations[index + i] = param;
+        }
 
         self.visit_stmt(arena, proc.name, &proc.body);
     }
@@ -193,14 +292,8 @@ impl<'a> ParamMap<'a> {
                     remainder: v,
                     body: b,
                 } => {
-                    let already_in_there = self
+                    self.join_points
-                        .items
+                        .insert(*j, Self::init_borrow_params(arena, xs));
-                        .insert(Key::JoinPoint(*j), Self::init_borrow_params(arena, xs));
-                    debug_assert!(
-                        already_in_there.is_none(),
-                        "join point {:?} is already defined!",
-                        j
-                    );
 
                     stack.push(v);
                     stack.push(b);
@@ -237,7 +330,6 @@ struct BorrowInfState<'a> {
     param_set: MutSet<Symbol>,
     owned: MutMap<Symbol, MutSet<Symbol>>,
     modified: bool,
-    param_map: ParamMap<'a>,
     arena: &'a Bump,
 }
 
@@ -245,14 +337,30 @@ impl<'a> BorrowInfState<'a> {
     pub fn own_var(&mut self, x: Symbol) {
         let current = self.owned.get_mut(&self.current_proc).unwrap();
 
-        if current.contains(&x) {
+        if current.insert(x) {
-            // do nothing
+            // entered if key was not yet present. If so, the set is modified,
-        } else {
+            // hence we set this flag
-            current.insert(x);
             self.modified = true;
         }
     }
 
+    /// if the extracted value is owned, then the surrounding structure must be too
+    fn if_is_owned_then_own(&mut self, extracted: Symbol, structure: Symbol) {
+        match self.owned.get_mut(&self.current_proc) {
+            None => unreachable!(
+                "the current procedure symbol {:?} is not in the owned map",
+                self.current_proc
+            ),
+            Some(set) => {
+                if set.contains(&extracted) && set.insert(structure) {
+                    // entered if key was not yet present. If so, the set is modified,
+                    // hence we set this flag
+                    self.modified = true;
+                }
+            }
+        }
+    }
+
     fn is_owned(&self, x: Symbol) -> bool {
         match self.owned.get(&self.current_proc) {
             None => unreachable!(
@@ -263,28 +371,50 @@ impl<'a> BorrowInfState<'a> {
         }
     }
 
-    fn update_param_map(&mut self, k: Key<'a>) {
+    fn update_param_map_help(&mut self, ps: &[Param<'a>]) -> &'a [Param<'a>] {
-        let arena = self.arena;
+        let mut new_ps = Vec::with_capacity_in(ps.len(), self.arena);
-        if let Some(ps) = self.param_map.items.get(&k) {
+        new_ps.extend(ps.iter().map(|p| {
-            let ps = Vec::from_iter_in(
-                ps.iter().map(|p| {
             if !p.borrow {
-                        p.clone()
+                *p
             } else if self.is_owned(p.symbol) {
                 self.modified = true;
-                        let mut p = p.clone();
+                let mut p = *p;
                 p.borrow = false;
 
                 p
             } else {
-                        p.clone()
+                *p
             }
-                }),
+        }));
-                arena,
-            );
 
-            self.param_map.items.insert(k, ps.into_bump_slice());
+        new_ps.into_bump_slice()
     }
+
+    fn update_param_map_declaration(
+        &mut self,
+        param_map: &mut ParamMap<'a>,
+        start: ParamOffset,
+        length: usize,
+    ) {
+        let index: usize = start.into();
+        let ps = &mut param_map.declarations[index..][..length];
+
+        for p in ps.iter_mut() {
+            if !p.borrow {
+                // do nothing
+            } else if self.is_owned(p.symbol) {
+                self.modified = true;
+                p.borrow = false;
+            } else {
+                // do nothing
+            }
+        }
+    }
+
+    fn update_param_map_join_point(&mut self, param_map: &mut ParamMap<'a>, id: JoinPointId) {
+        let ps = param_map.join_points[&id];
+        let new_ps = self.update_param_map_help(ps);
+        param_map.join_points.insert(id, new_ps);
     }
 
     /// This looks at an application `f x1 x2 x3`
@@ -351,7 +481,7 @@ impl<'a> BorrowInfState<'a> {
     ///
     /// and determines whether z and which of the symbols used in e
     /// must be taken as owned parameters
-    fn collect_call(&mut self, z: Symbol, e: &crate::ir::Call<'a>) {
+    fn collect_call(&mut self, param_map: &mut ParamMap<'a>, z: Symbol, e: &crate::ir::Call<'a>) {
         use crate::ir::CallType::*;
 
         let crate::ir::Call {
@@ -369,8 +499,7 @@ impl<'a> BorrowInfState<'a> {
                 let top_level = ProcLayout::new(self.arena, arg_layouts, *ret_layout);
 
                 // get the borrow signature of the applied function
-                let ps = self
+                let ps = param_map
-                    .param_map
                     .get_symbol(*name, top_level)
                     .expect("function is defined");
 
@@ -386,6 +515,7 @@ impl<'a> BorrowInfState<'a> {
                     ps.len(),
                     arguments.len()
                 );
+
                 self.own_args_using_params(arguments, ps);
             }
 
@@ -416,7 +546,7 @@ impl<'a> BorrowInfState<'a> {
 
                 match op {
                     ListMap | ListKeepIf | ListKeepOks | ListKeepErrs => {
-                        match self.param_map.get_symbol(arguments[1], closure_layout) {
+                        match param_map.get_symbol(arguments[1], closure_layout) {
                             Some(function_ps) => {
                                 // own the list if the function wants to own the element
                                 if !function_ps[0].borrow {
@@ -432,7 +562,7 @@ impl<'a> BorrowInfState<'a> {
                         }
                     }
                     ListMapWithIndex => {
-                        match self.param_map.get_symbol(arguments[1], closure_layout) {
+                        match param_map.get_symbol(arguments[1], closure_layout) {
                             Some(function_ps) => {
                                 // own the list if the function wants to own the element
                                 if !function_ps[1].borrow {
@@ -447,7 +577,7 @@ impl<'a> BorrowInfState<'a> {
                             None => unreachable!(),
                         }
                     }
-                    ListMap2 => match self.param_map.get_symbol(arguments[2], closure_layout) {
+                    ListMap2 => match param_map.get_symbol(arguments[2], closure_layout) {
                         Some(function_ps) => {
                             // own the lists if the function wants to own the element
                             if !function_ps[0].borrow {
@@ -465,7 +595,7 @@ impl<'a> BorrowInfState<'a> {
                         }
                         None => unreachable!(),
                     },
-                    ListMap3 => match self.param_map.get_symbol(arguments[3], closure_layout) {
+                    ListMap3 => match param_map.get_symbol(arguments[3], closure_layout) {
                         Some(function_ps) => {
                             // own the lists if the function wants to own the element
                             if !function_ps[0].borrow {
@@ -486,7 +616,7 @@ impl<'a> BorrowInfState<'a> {
                         None => unreachable!(),
                     },
                     ListSortWith => {
-                        match self.param_map.get_symbol(arguments[1], closure_layout) {
+                        match param_map.get_symbol(arguments[1], closure_layout) {
                             Some(function_ps) => {
                                 // always own the input list
                                 self.own_var(arguments[0]);
@@ -500,7 +630,7 @@ impl<'a> BorrowInfState<'a> {
                         }
                     }
                     ListWalk | ListWalkUntil | ListWalkBackwards | DictWalk => {
-                        match self.param_map.get_symbol(arguments[2], closure_layout) {
+                        match param_map.get_symbol(arguments[2], closure_layout) {
                             Some(function_ps) => {
                                 // own the data structure if the function wants to own the element
                                 if !function_ps[0].borrow {
@@ -542,7 +672,7 @@ impl<'a> BorrowInfState<'a> {
         }
     }
 
-    fn collect_expr(&mut self, z: Symbol, e: &Expr<'a>) {
+    fn collect_expr(&mut self, param_map: &mut ParamMap<'a>, z: Symbol, e: &Expr<'a>) {
         use Expr::*;
 
         match e {
@@ -570,50 +700,44 @@ impl<'a> BorrowInfState<'a> {
                 self.own_var(z);
             }
 
-            Call(call) => self.collect_call(z, call),
+            Call(call) => self.collect_call(param_map, z, call),
 
             Literal(_) | RuntimeErrorFunction(_) => {}
 
             StructAtIndex { structure: x, .. } => {
                 // if the structure (record/tag/array) is owned, the extracted value is
-                if self.is_owned(*x) {
+                self.if_is_owned_then_own(*x, z);
-                    self.own_var(z);
-                }
 
                 // if the extracted value is owned, the structure must be too
-                if self.is_owned(z) {
+                self.if_is_owned_then_own(z, *x);
-                    self.own_var(*x);
-                }
             }
 
             UnionAtIndex { structure: x, .. } => {
                 // if the structure (record/tag/array) is owned, the extracted value is
-                if self.is_owned(*x) {
+                self.if_is_owned_then_own(*x, z);
-                    self.own_var(z);
-                }
 
                 // if the extracted value is owned, the structure must be too
-                if self.is_owned(z) {
+                self.if_is_owned_then_own(z, *x);
-                    self.own_var(*x);
-                }
             }
 
             GetTagId { structure: x, .. } => {
                 // if the structure (record/tag/array) is owned, the extracted value is
-                if self.is_owned(*x) {
+                self.if_is_owned_then_own(*x, z);
-                    self.own_var(z);
-                }
 
                 // if the extracted value is owned, the structure must be too
-                if self.is_owned(z) {
+                self.if_is_owned_then_own(z, *x);
-                    self.own_var(*x);
-                }
             }
         }
     }
 
     #[allow(clippy::many_single_char_names)]
-    fn preserve_tail_call(&mut self, x: Symbol, v: &Expr<'a>, b: &Stmt<'a>) {
+    fn preserve_tail_call(
+        &mut self,
+        param_map: &mut ParamMap<'a>,
+        x: Symbol,
+        v: &Expr<'a>,
+        b: &Stmt<'a>,
+    ) {
         if let (
             Expr::Call(crate::ir::Call {
                 call_type:
@@ -634,7 +758,7 @@ impl<'a> BorrowInfState<'a> {
             if self.current_proc == *g && x == *z {
                 // anonymous functions (for which the ps may not be known)
                 // can never be tail-recursive, so this is fine
-                if let Some(ps) = self.param_map.get_symbol(*g, top_level) {
+                if let Some(ps) = param_map.get_symbol(*g, top_level) {
                     self.own_params_using_args(ys, ps)
                 }
             }
@@ -653,7 +777,7 @@ impl<'a> BorrowInfState<'a> {
         }
     }
 
-    fn collect_stmt(&mut self, stmt: &Stmt<'a>) {
+    fn collect_stmt(&mut self, param_map: &mut ParamMap<'a>, stmt: &Stmt<'a>) {
         use Stmt::*;
 
         match stmt {
@@ -665,17 +789,35 @@ impl<'a> BorrowInfState<'a> {
             } => {
                 let old = self.param_set.clone();
                 self.update_param_set(ys);
-                self.collect_stmt(v);
+                self.collect_stmt(param_map, v);
                 self.param_set = old;
-                self.update_param_map(Key::JoinPoint(*j));
+                self.update_param_map_join_point(param_map, *j);
 
-                self.collect_stmt(b);
+                self.collect_stmt(param_map, b);
             }
 
-            Let(x, v, _, b) => {
+            Let(x, v, _, mut b) => {
-                self.collect_stmt(b);
+                let mut stack = Vec::new_in(self.arena);
-                self.collect_expr(*x, v);
+
-                self.preserve_tail_call(*x, v, b);
+                stack.push((*x, v));
+
+                while let Stmt::Let(symbol, expr, _, tail) = b {
+                    b = tail;
+                    stack.push((*symbol, expr));
+                }
+
+                self.collect_stmt(param_map, b);
+
+                let mut it = stack.into_iter().rev();
+
+                // collect the final expr, and see if we need to preserve a tail call
+                let (x, v) = it.next().unwrap();
+                self.collect_expr(param_map, x, v);
+                self.preserve_tail_call(param_map, x, v, b);
+
+                for (x, v) in it {
+                    self.collect_expr(param_map, x, v);
+                }
             }
 
             Invoke {
@@ -686,17 +828,17 @@ impl<'a> BorrowInfState<'a> {
                 fail,
                 exception_id: _,
             } => {
-                self.collect_stmt(pass);
+                self.collect_stmt(param_map, pass);
-                self.collect_stmt(fail);
+                self.collect_stmt(param_map, fail);
 
-                self.collect_call(*symbol, call);
+                self.collect_call(param_map, *symbol, call);
 
                 // TODO how to preserve the tail call of an invoke?
                 // self.preserve_tail_call(*x, v, b);
             }
 
             Jump(j, ys) => {
-                let ps = self.param_map.get_join_point(*j);
+                let ps = param_map.get_join_point(*j);
 
                 // for making sure the join point can reuse
                 self.own_args_using_params(ys, ps);
@@ -710,9 +852,9 @@ impl<'a> BorrowInfState<'a> {
                 ..
             } => {
                 for (_, _, b) in branches.iter() {
-                    self.collect_stmt(b);
+                    self.collect_stmt(param_map, b);
                 }
-                self.collect_stmt(default_branch.1);
+                self.collect_stmt(param_map, default_branch.1);
             }
             Refcounting(_, _) => unreachable!("these have not been introduced yet"),
 
@@ -722,7 +864,12 @@ impl<'a> BorrowInfState<'a> {
         }
     }
 
-    fn collect_proc(&mut self, proc: &Proc<'a>, layout: ProcLayout<'a>) {
+    fn collect_proc(
+        &mut self,
+        param_map: &mut ParamMap<'a>,
+        proc: &Proc<'a>,
+        param_offset: ParamOffset,
+    ) {
         let old = self.param_set.clone();
 
         let ys = Vec::from_iter_in(proc.args.iter().map(|t| t.1), self.arena).into_bump_slice();
@@ -732,8 +879,8 @@ impl<'a> BorrowInfState<'a> {
         // ensure that current_proc is in the owned map
         self.owned.entry(proc.name).or_default();
 
-        self.collect_stmt(&proc.body);
+        self.collect_stmt(param_map, &proc.body);
-        self.update_param_map(Key::Declaration(proc.name, layout));
+        self.update_param_map_declaration(param_map, param_offset, proc.args.len());
 
         self.param_set = old;
     }
@@ -827,3 +974,87 @@ pub fn lowlevel_borrow_signature(arena: &Bump, op: LowLevel) -> &[bool] {
         ExpectTrue => arena.alloc_slice_copy(&[irrelevant]),
     }
 }
+
+fn make_successor_mapping<'a>(
+    arena: &'a Bump,
+    procs: &MutMap<(Symbol, ProcLayout<'_>), Proc<'a>>,
+) -> MutMap<Symbol, Vec<'a, Symbol>> {
+    let mut result = MutMap::with_capacity_and_hasher(procs.len(), default_hasher());
+
+    for (key, proc) in procs {
+        let mut call_info = CallInfo {
+            keys: Vec::new_in(arena),
+        };
+        call_info_stmt(arena, &proc.body, &mut call_info);
+
+        let mut keys = call_info.keys;
+        keys.sort_unstable();
+        keys.dedup();
+
+        result.insert(key.0, keys);
+    }
+
+    result
+}
+
+struct CallInfo<'a> {
+    keys: Vec<'a, Symbol>,
+}
+
+fn call_info_call<'a>(call: &crate::ir::Call<'a>, info: &mut CallInfo<'a>) {
+    use crate::ir::CallType::*;
+
+    match call.call_type {
+        ByName { name, .. } => {
+            info.keys.push(name);
+        }
+        Foreign { .. } => {}
+        LowLevel { .. } => {}
+        HigherOrderLowLevel { .. } => {}
+    }
+}
+
+fn call_info_stmt<'a>(arena: &'a Bump, stmt: &Stmt<'a>, info: &mut CallInfo<'a>) {
+    use Stmt::*;
+
+    let mut stack = bumpalo::vec![ in arena; stmt ];
+
+    while let Some(stmt) = stack.pop() {
+        match stmt {
+            Join {
+                remainder: v,
+                body: b,
+                ..
+            } => {
+                stack.push(v);
+                stack.push(b);
+            }
+            Let(_, expr, _, cont) => {
+                if let Expr::Call(call) = expr {
+                    call_info_call(call, info);
+                }
+                stack.push(cont);
+            }
+            Invoke {
+                call, pass, fail, ..
+            } => {
+                call_info_call(call, info);
+                stack.push(pass);
+                stack.push(fail);
+            }
+            Switch {
+                branches,
+                default_branch,
+                ..
+            } => {
+                stack.extend(branches.iter().map(|b| &b.2));
+                stack.push(default_branch.1);
+            }
+            Refcounting(_, _) => unreachable!("these have not been introduced yet"),
+
+            Ret(_) | Resume(_) | Jump(_, _) | RuntimeError(_) => {
+                // these are terminal, do nothing
+            }
+        }
+    }
+}

compiler/mono/src/inc_dec.rs

@@ -247,8 +247,7 @@ impl<'a> Context<'a> {
     pub fn new(arena: &'a Bump, param_map: &'a ParamMap<'a>) -> Self {
         let mut vars = MutMap::default();
 
-        for (key, _) in param_map.into_iter() {
+        for symbol in param_map.iter_symbols() {
-            if let crate::borrow::Key::Declaration(symbol, _) = key {
             vars.insert(
                 *symbol,
                 VarInfo {
@@ -259,7 +258,6 @@ impl<'a> Context<'a> {
                 },
             );
         }
-        }
 
         Self {
             arena,
@@ -1261,28 +1259,25 @@ fn update_jp_live_vars(j: JoinPointId, ys: &[Param], v: &Stmt<'_>, m: &mut JPLiv
     m.insert(j, j_live_vars);
 }
 
-/// used to process the main function in the repl
+pub fn visit_procs<'a>(
-pub fn visit_declaration<'a>(
     arena: &'a Bump,
     param_map: &'a ParamMap<'a>,
-    stmt: &'a Stmt<'a>,
+    procs: &mut MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>,
-) -> &'a Stmt<'a> {
-    let ctx = Context::new(arena, param_map);
-
-    let params = &[] as &[_];
-    let ctx = ctx.update_var_info_with_params(params);
-    let (b, b_live_vars) = ctx.visit_stmt(stmt);
-    ctx.add_dec_for_dead_params(params, b, &b_live_vars)
-}
-
-pub fn visit_proc<'a>(
-    arena: &'a Bump,
-    param_map: &'a ParamMap<'a>,
-    proc: &mut Proc<'a>,
-    layout: ProcLayout<'a>,
 ) {
     let ctx = Context::new(arena, param_map);
 
+    for (key, proc) in procs.iter_mut() {
+        visit_proc(arena, param_map, &ctx, proc, key.1);
+    }
+}
+
+fn visit_proc<'a>(
+    arena: &'a Bump,
+    param_map: &'a ParamMap<'a>,
+    ctx: &Context<'a>,
+    proc: &mut Proc<'a>,
+    layout: ProcLayout<'a>,
+) {
     let params = match param_map.get_symbol(proc.name, layout) {
         Some(slice) => slice,
         None => Vec::from_iter_in(

compiler/mono/src/ir.rs

@@ -222,9 +222,7 @@ impl<'a> Proc<'a> {
     ) {
         let borrow_params = arena.alloc(crate::borrow::infer_borrow(arena, procs));
 
-        for (key, proc) in procs.iter_mut() {
+        crate::inc_dec::visit_procs(arena, borrow_params, procs);
-            crate::inc_dec::visit_proc(arena, borrow_params, proc, key.1);
-        }
     }
 
     pub fn insert_reset_reuse_operations<'i>(
@@ -430,9 +428,7 @@ impl<'a> Procs<'a> {
 
         let borrow_params = arena.alloc(crate::borrow::infer_borrow(arena, &result));
 
-        for (key, proc) in result.iter_mut() {
+        crate::inc_dec::visit_procs(arena, borrow_params, &mut result);
-            crate::inc_dec::visit_proc(arena, borrow_params, proc, key.1);
-        }
 
         result
     }
@@ -473,9 +469,7 @@ impl<'a> Procs<'a> {
 
         let borrow_params = arena.alloc(crate::borrow::infer_borrow(arena, &result));
 
-        for (key, proc) in result.iter_mut() {
+        crate::inc_dec::visit_procs(arena, borrow_params, &mut result);
-            crate::inc_dec::visit_proc(arena, borrow_params, proc, key.1);
-        }
 
         (result, borrow_params)
     }
@@ -848,13 +842,21 @@ impl<'a, 'i> Env<'a, 'i> {
 #[derive(Clone, Debug, PartialEq, Copy, Eq, Hash)]
 pub struct JoinPointId(pub Symbol);
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub struct Param<'a> {
     pub symbol: Symbol,
     pub borrow: bool,
     pub layout: Layout<'a>,
 }
 
+impl<'a> Param<'a> {
+    pub const EMPTY: Self = Param {
+        symbol: Symbol::EMPTY_PARAM,
+        borrow: false,
+        layout: Layout::Struct(&[]),
+    };
+}
+
 pub fn cond<'a>(
     env: &mut Env<'a, '_>,
     cond_symbol: Symbol,