diff --git a/compiler/mono/src/ir.rs b/compiler/mono/src/ir.rs index b39bb8b722..b0635bbf1f 100644 --- a/compiler/mono/src/ir.rs +++ b/compiler/mono/src/ir.rs @@ -272,6 +272,33 @@ impl<'a> Proc<'a> { proc.body = b.clone(); } } + + fn make_tail_recursive(&mut self, env: &mut Env<'a, '_>) { + let mut args = Vec::with_capacity_in(self.args.len(), env.arena); + let mut proc_args = Vec::with_capacity_in(self.args.len(), env.arena); + + for (layout, symbol) in self.args { + let new = env.unique_symbol(); + args.push((*layout, *symbol, new)); + proc_args.push((*layout, new)); + } + + use self::SelfRecursive::*; + if let SelfRecursive(id) = self.is_self_recursive { + let transformed = crate::tail_recursion::make_tail_recursive( + env.arena, + id, + self.name, + self.body.clone(), + args.into_bump_slice(), + ); + + if let Some(with_tco) = transformed { + self.body = with_tco; + self.args = proc_args.into_bump_slice(); + } + } + } } #[derive(Clone, Debug)] @@ -350,7 +377,7 @@ pub enum InProgressProc<'a> { impl<'a> Procs<'a> { pub fn get_specialized_procs_without_rc( self, - arena: &'a Bump, + env: &mut Env<'a, '_>, ) -> MutMap<(Symbol, ProcLayout<'a>), Proc<'a>> { let mut result = MutMap::with_capacity_and_hasher(self.specialized.len(), default_hasher()); @@ -376,16 +403,7 @@ impl<'a> Procs<'a> { panic!(); } Done(mut proc) => { - use self::SelfRecursive::*; - if let SelfRecursive(id) = proc.is_self_recursive { - proc.body = crate::tail_recursion::make_tail_recursive( - arena, - id, - proc.name, - proc.body.clone(), - proc.args, - ); - } + proc.make_tail_recursive(env); result.insert(key, proc); } @@ -395,86 +413,6 @@ impl<'a> Procs<'a> { result } - // TODO investigate make this an iterator? - pub fn get_specialized_procs( - self, - arena: &'a Bump, - ) -> MutMap<(Symbol, ProcLayout<'a>), Proc<'a>> { - let mut result = MutMap::with_capacity_and_hasher(self.specialized.len(), default_hasher()); - - for ((s, toplevel), in_prog_proc) in self.specialized.into_iter() { - match in_prog_proc { - InProgress => unreachable!( - "The procedure {:?} should have be done by now", - (s, toplevel) - ), - Done(proc) => { - result.insert((s, toplevel), proc); - } - } - } - - for (_, proc) in result.iter_mut() { - use self::SelfRecursive::*; - if let SelfRecursive(id) = proc.is_self_recursive { - proc.body = crate::tail_recursion::make_tail_recursive( - arena, - id, - proc.name, - proc.body.clone(), - proc.args, - ); - } - } - - let borrow_params = arena.alloc(crate::borrow::infer_borrow(arena, &result)); - - crate::inc_dec::visit_procs(arena, borrow_params, &mut result); - - result - } - - pub fn get_specialized_procs_help( - self, - arena: &'a Bump, - ) -> ( - MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>, - &'a crate::borrow::ParamMap<'a>, - ) { - let mut result = MutMap::with_capacity_and_hasher(self.specialized.len(), default_hasher()); - - for ((s, toplevel), in_prog_proc) in self.specialized.into_iter() { - match in_prog_proc { - InProgress => unreachable!( - "The procedure {:?} should have be done by now", - (s, toplevel) - ), - Done(proc) => { - result.insert((s, toplevel), proc); - } - } - } - - for (_, proc) in result.iter_mut() { - use self::SelfRecursive::*; - if let SelfRecursive(id) = proc.is_self_recursive { - proc.body = crate::tail_recursion::make_tail_recursive( - arena, - id, - proc.name, - proc.body.clone(), - proc.args, - ); - } - } - - let borrow_params = arena.alloc(crate::borrow::infer_borrow(arena, &result)); - - crate::inc_dec::visit_procs(arena, borrow_params, &mut result); - - (result, borrow_params) - } - // TODO trim down these arguments! #[allow(clippy::too_many_arguments)] pub fn insert_named( diff --git a/compiler/mono/src/tail_recursion.rs b/compiler/mono/src/tail_recursion.rs index d0a75b0c2a..08dc545fc3 100644 --- a/compiler/mono/src/tail_recursion.rs +++ b/compiler/mono/src/tail_recursion.rs @@ -33,16 +33,16 @@ pub fn make_tail_recursive<'a>( id: JoinPointId, needle: Symbol, stmt: Stmt<'a>, - args: &'a [(Layout<'a>, Symbol)], -) -> Stmt<'a> { + args: &'a [(Layout<'a>, Symbol, Symbol)], +) -> Option> { let allocated = arena.alloc(stmt); match insert_jumps(arena, allocated, id, needle) { - None => allocated.clone(), + None => None, Some(new) => { // jumps were inserted, we must now add a join point let params = Vec::from_iter_in( - args.iter().map(|(layout, symbol)| Param { + args.iter().map(|(layout, symbol, _)| Param { symbol: *symbol, layout: *layout, borrow: true, @@ -52,16 +52,18 @@ pub fn make_tail_recursive<'a>( .into_bump_slice(); // TODO could this be &[]? - let args = Vec::from_iter_in(args.iter().map(|t| t.1), arena).into_bump_slice(); + let args = Vec::from_iter_in(args.iter().map(|t| t.2), arena).into_bump_slice(); let jump = arena.alloc(Stmt::Jump(id, args)); - Stmt::Join { + let join = Stmt::Join { id, remainder: jump, parameters: params, body: new, - } + }; + + Some(join) } } }