diff --git a/crates/compiler/mono/src/fl_reuse.rs b/crates/compiler/mono/src/fl_reuse.rs index 70af9ad491..afceec4e0b 100644 --- a/crates/compiler/mono/src/fl_reuse.rs +++ b/crates/compiler/mono/src/fl_reuse.rs @@ -1,11 +1,12 @@ // Frame limited reuse // Based on Reference Counting with Frame Limited Reuse -use crate::ir::{BranchInfo, Expr, ModifyRc, Proc, ProcLayout, Stmt, UpdateModeIds}; +use crate::ir::{BranchInfo, Expr, ModifyRc, Proc, ProcLayout, Stmt, UpdateModeId, UpdateModeIds}; use crate::layout::{InLayout, Layout, LayoutInterner, STLayoutInterner, UnionLayout}; use bumpalo::Bump; +use bumpalo::collections::vec::Vec; use roc_collections::MutMap; use roc_module::symbol::{IdentIds, ModuleId, Symbol}; @@ -21,6 +22,11 @@ pub fn insert_reset_reuse_operations<'a, 'i>( update_mode_ids: &'i mut UpdateModeIds, procs: &mut MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>, ) { + let mut global_layouts = SymbolLayout::default(); + for (symbol, layout) in procs.keys() { + global_layouts.insert(*symbol, LayoutOption::GloballyDefined); + } + for proc in procs.values_mut() { let new_proc = insert_reset_reuse_operations_proc( arena, @@ -28,21 +34,31 @@ pub fn insert_reset_reuse_operations<'a, 'i>( home, ident_ids, update_mode_ids, + global_layouts.clone(), proc.clone(), ); *proc = new_proc; } } -pub fn insert_reset_reuse_operations_proc<'a, 'i>( +fn insert_reset_reuse_operations_proc<'a, 'i>( arena: &'a Bump, layout_interner: &'i mut STLayoutInterner<'a>, home: ModuleId, ident_ids: &'i mut IdentIds, update_mode_ids: &'i mut UpdateModeIds, + mut symbol_layout: SymbolLayout<'a>, mut proc: Proc<'a>, ) -> Proc<'a> { - let env = ReuseEnvironment::default(); + for (layout, symbol) in proc.args { + symbol_layout.insert(*symbol, LayoutOption::Layout(layout)); + } + + let mut env = ReuseEnvironment { + layout_tags: MutMap::default(), + reuse_tokens: MutMap::default(), + symbol_layouts: symbol_layout, + }; let new_body = insert_reset_reuse_operations_stmt( arena, @@ -50,10 +66,13 @@ pub fn insert_reset_reuse_operations_proc<'a, 'i>( home, ident_ids, update_mode_ids, - env, + &mut env, arena.alloc(proc.body), ); + // All reuse tokens either have to be used by reuse or not be inserted at all at the reset (and removed from the environment). + debug_assert!(env.reuse_tokens.is_empty()); + proc.body = new_body.clone(); proc } @@ -64,22 +83,41 @@ fn insert_reset_reuse_operations_stmt<'a, 'i>( home: ModuleId, ident_ids: &'i mut IdentIds, update_mode_ids: &'i mut UpdateModeIds, - mut environment: ReuseEnvironment<'a>, + environment: &mut ReuseEnvironment<'a>, stmt: &'a Stmt<'a>, ) -> &'a Stmt<'a> { match stmt { Stmt::Let(binding, expr, layout, continuation) => { - // TODO store the layout for the binding in the environment. - // TODO deal with fresh allocations w/ reuse. - if let Expr::Tag { - tag_layout, - tag_id, - arguments, - } = expr - { - // The value of the tag is currently only used in the case of nullable recursive unions. - // But for completeness we add every kind of union to the layout_tags. - environment.add_layout_tag(layout, *tag_id); + let new_expr = match expr { + Expr::Tag { + tag_layout, + tag_id, + arguments, + } => { + // The value of the tag is currently only used in the case of nullable recursive unions. + // But for completeness we add every kind of union to the layout_tags. + environment.add_layout_tag(layout, *tag_id); + + // See if we have a reuse token + match environment.pop_reuse_token(layout) { + // We have a reuse token for this layout, use it. + Some(reuse_token) => { + Expr::Reuse { + symbol: reuse_token.symbol, + update_mode: reuse_token.update_mode_id, + // for now, always overwrite the tag ID just to be sure + update_tag_id: true, + tag_layout: *tag_layout, + tag_id: *tag_id, + arguments, + } + } + + // We have no reuse token available, keep the old expression with a fresh allocation. + None => expr.clone(), + } + } + _ => expr.clone(), }; environment.add_symbol_layout(*binding, layout); @@ -94,11 +132,14 @@ fn insert_reset_reuse_operations_stmt<'a, 'i>( continuation, ); - arena.alloc(Stmt::Let(*binding, expr.clone(), *layout, new_continuation)) + // // for now, always overwrite the tag ID just to be sure + // let update_tag_id = true; + + arena.alloc(Stmt::Let(*binding, new_expr, *layout, new_continuation)) } Stmt::Switch { cond_symbol, - cond_layout: layout, + cond_layout, branches, default_branch, ret_layout, @@ -108,7 +149,7 @@ fn insert_reset_reuse_operations_stmt<'a, 'i>( .map(|(tag_id, info, branch)| { let mut branch_env = environment.clone(); if let BranchInfo::Constructor { tag_id: tag, .. } = info { - branch_env.add_layout_tag(layout, *tag); + branch_env.add_layout_tag(cond_layout, *tag); } let new_branch = insert_reset_reuse_operations_stmt( @@ -117,11 +158,11 @@ fn insert_reset_reuse_operations_stmt<'a, 'i>( home, ident_ids, update_mode_ids, - environment.clone(), + &mut environment.clone(), branch, ); - (*tag_id, info, new_branch, branch_env) + (*tag_id, info.clone(), new_branch, branch_env) }) .collect::>(); @@ -130,7 +171,7 @@ fn insert_reset_reuse_operations_stmt<'a, 'i>( let mut branch_env = environment.clone(); if let BranchInfo::Constructor { tag_id: tag, .. } = info { - branch_env.add_layout_tag(layout, *tag); + branch_env.add_layout_tag(cond_layout, *tag); } let new_branch = insert_reset_reuse_operations_stmt( @@ -139,40 +180,126 @@ fn insert_reset_reuse_operations_stmt<'a, 'i>( home, ident_ids, update_mode_ids, - environment.clone(), + &mut environment.clone(), branch, ); (info.clone(), new_branch, branch_env) }; - // TODO collect the used reuse token from each branch and drop the unused ones. - // Make sure the current env is updated to have the correct reuse tokens consumed. + // First we determine the minimum of reuse tokens available for each layout. + let layout_min_reuse_tokens = { + let branch_envs = { + let mut branch_environments = + Vec::with_capacity_in(new_branches.len() + 1, arena); - // TODO add value to layout_tags - todo!() - } - Stmt::Refcounting(rc, continuation) => { - match rc { - ModifyRc::Inc(_, _) => todo!(), - // TODO verify whether reuse works for a decref as well. It should be. - ModifyRc::Dec(symbol) | ModifyRc::DecRef(symbol) => { - // TODO insert actual reuse statements. if the reuse token is consumed (check by checking the top of the stack) + for (_, _, _, branch_env) in new_branches.iter() { + branch_environments.push(branch_env); + } - // Get the layout of the symbol from where it is defined. - let layout = environment.get_symbol_layout(*symbol); + branch_environments.push(&new_default_branch.2); - // If the layout is reusable, add a reuse token for the layout to the environment. - if matches!( - can_reuse_layout_tag(layout_interner, &environment, &layout), - Reuse::Reusable - ) { - environment.push_reuse_token(&layout, *symbol); - }; - } + branch_environments + }; + + let layout_min_reuse_tokens = + environment.reuse_tokens.keys().copied().map(|layout| { + let min_reuse_tokens = branch_envs + .iter() + .map(|branch_environment| { + branch_environment + .reuse_tokens + .get(&layout) + .map_or(0, |tokens| tokens.len()) + }) + .min() + .expect("There should be at least one branch"); + (layout, min_reuse_tokens) + }); + + layout_min_reuse_tokens.collect::>() + }; + + // Then we drop any unused reuse tokens in branches where the minimum is not reached. + let newer_branches = Vec::from_iter_in( + new_branches + .into_iter() + .map(|(label, info, branch, branch_env)| { + let unused_tokens= branch_env.reuse_tokens.iter().flat_map(|(layout, reuse_tokens) |{ + let min_reuse_tokens = layout_min_reuse_tokens.get(&layout).expect("All layouts in the environment should be in the layout_min_reuse_tokens map."); + let unused_tokens = &reuse_tokens[*min_reuse_tokens..]; + unused_tokens + }); + + let newer_branch = drop_unused_reuse_tokens(arena, unused_tokens.copied(), branch); + + (label, info, newer_branch.clone()) + }), + arena, + ) + .into_bump_slice(); + + let newer_default_branch = { + let (info, branch, branch_env) = new_default_branch; + let unused_tokens= branch_env.reuse_tokens.iter().flat_map(|(layout, reuse_tokens) |{ + let min_reuse_tokens = layout_min_reuse_tokens.get(&layout).expect("All layouts in the environment should be in the layout_min_reuse_tokens map."); + let unused_tokens = &reuse_tokens[*min_reuse_tokens..]; + unused_tokens + }); + + let newer_branch = drop_unused_reuse_tokens(arena, unused_tokens.copied(), branch); + + (info, newer_branch) + }; + + // And finally we update the current environment to reflect the correct number of reuse tokens. + for (layout, reuse_tokens) in environment.reuse_tokens.iter_mut() { + let min_reuse_tokens = layout_min_reuse_tokens.get(&layout).expect( + "All layouts in the environment should be in the layout_min_reuse_tokens map.", + ); + reuse_tokens.truncate(*min_reuse_tokens) } - insert_reset_reuse_operations_stmt( + arena.alloc(Stmt::Switch { + cond_symbol: *cond_symbol, + cond_layout: *cond_layout, + branches: newer_branches, + default_branch: newer_default_branch, + ret_layout: *ret_layout, + }) + } + Stmt::Refcounting(rc, continuation) => { + let reuse_pair = match rc { + ModifyRc::Inc(_, _) => { + // We don't need to do anything for an inc. + None + } + ModifyRc::Dec(symbol) | ModifyRc::DecRef(symbol) => { + // Get the layout of the symbol from where it is defined. + let layout_option = environment.get_symbol_layout(*symbol); + + // If the symbol is defined in the current proc, we can use the layout from the environment. + match layout_option { + LayoutOption::Layout(layout) + if matches!( + can_reuse_layout_tag(layout_interner, &environment, &layout), + Reuse::Reusable + ) => + { + let layout_clone = layout.clone(); + let reuse_token = ReuseToken { + symbol: Symbol::new(home, ident_ids.gen_unique()), + update_mode_id: update_mode_ids.next_id(), + }; + environment.push_reuse_token(&layout_clone, reuse_token); + Some((layout_clone, *symbol, reuse_token)) + } + _ => None, + } + } + }; + + let new_continuation = insert_reset_reuse_operations_stmt( arena, layout_interner, home, @@ -180,11 +307,43 @@ fn insert_reset_reuse_operations_stmt<'a, 'i>( update_mode_ids, environment, continuation, - ) + ); + + // If we inserted a reuse token, we need to insert a reset reuse operation if the reuse token is consumed. + if let Some((layout, symbol, reuse_token)) = reuse_pair { + let stack_reuse_token = environment.peek_reuse_token(&layout); + + match stack_reuse_token { + Some(reuse_symbol) if reuse_symbol == reuse_token => { + // The token we inserted is still on the stack, so we don't need to insert a reset operation. + // We do need to remove the token from the environment. To prevent errors higher in the tree. + let _ = environment.pop_reuse_token(&layout); + } + _ => { + // The token we inserted is no longer on the stack, it must have been consumed. + // So we need to insert a reset operation. + let reset_expr = Expr::Reset { + symbol, + update_mode: reuse_token.update_mode_id, + }; + + // If we generate a reuse token, we no longer want to use the drop statement anymore. So we just return the reset expression. + // TODO verify if this works for both dec and decref. + return arena.alloc(Stmt::Let( + reuse_token.symbol, + reset_expr, + // TODO not sure what the layout should be for a reset token. Currently it is the layout of the symbol. + *layout, + new_continuation, + )); + } + } + } + + arena.alloc(Stmt::Refcounting(rc.clone(), new_continuation)) } - Stmt::Ret(symbol) => - // The return statement just doesn't consume any tokens. Dropping these tokens will be handled before. - { + Stmt::Ret(symbol) => { + // The return statement just doesn't consume any tokens. Dropping these tokens will be handled before. stmt } Stmt::Expect { @@ -289,27 +448,49 @@ enum Reuse { Map containing the curren't known tag of a layout. A layout with a tag will be inserted e.g. after pattern matching. where the tag is known. */ -type LayoutTags<'a> = MutMap<&'a InLayout<'a>, FooBar>; +type LayoutTags<'a> = MutMap<&'a InLayout<'a>, Tag>; /** Map containing the reuse tokens of a layout. A vec is used as a stack as we want to use the latest reuse token available. */ -type ReuseTokens<'a> = MutMap, Vec>; +type ReuseTokens<'a> = MutMap, std::vec::Vec>; /** A reuse token is a symbol that is used to reset a layout. Matches variables that are pointers. */ -type ResetToken = Symbol; +#[derive(Clone, Copy, PartialEq, Eq)] +struct ReuseToken { + // The symbol of the reuse token. + symbol: Symbol, -type FooBar = u16; + // TODO figure out what this is. + update_mode_id: UpdateModeId, +} + +type Tag = u16; + +/** +Map containing the layout of a symbol. +Used to determine whether the pointer of a symbol can be reused, if it is reference counted (heap allocated). +*/ +type SymbolLayout<'a> = MutMap>; + +#[derive(Clone)] +enum LayoutOption<'a> { + // A normal layout defined in the current function. + Layout(&'a InLayout<'a>), + + // No layout as this symbol is defined in a global scope and should not be reused. + GloballyDefined, +} #[derive(Default, Clone)] struct ReuseEnvironment<'a> { layout_tags: LayoutTags<'a>, reuse_tokens: ReuseTokens<'a>, - symbol_layouts: MutMap>, + symbol_layouts: SymbolLayout<'a>, } impl<'a> ReuseEnvironment<'a> { @@ -317,21 +498,21 @@ impl<'a> ReuseEnvironment<'a> { Add the known tag for a layout. Used to optimize reuse of unions that are know to have a null pointer. */ - fn add_layout_tag(&mut self, layout: &'a InLayout<'a>, tag: FooBar) { + fn add_layout_tag(&mut self, layout: &'a InLayout<'a>, tag: Tag) { self.layout_tags.insert(layout, tag); } /** Retrieve the known tag for a layout. */ - fn get_layout_tag(&self, layout: &InLayout<'a>) -> Option { + fn get_layout_tag(&self, layout: &InLayout<'a>) -> Option { self.layout_tags.get(layout).copied() } /** Retrieve a reuse token for a layout from the stack for said layout. */ - fn get_reuse_token(&mut self, layout: &InLayout<'a>) -> Option { + fn pop_reuse_token(&mut self, layout: &InLayout<'a>) -> Option { let reuse_tokens = self.reuse_tokens.get_mut(layout)?; // If the layout is in the map, pop the token from the stack. let reuse_token = reuse_tokens.pop(); @@ -342,10 +523,21 @@ impl<'a> ReuseEnvironment<'a> { reuse_token } + /** + Retrieve a reuse token for a layout from the stack for said layout. + Without consuming the token. + */ + fn peek_reuse_token(&mut self, layout: &InLayout<'a>) -> Option { + let reuse_tokens = self.reuse_tokens.get(layout)?; + // If the layout is in the map, pop the token from the stack. + let reuse_token = reuse_tokens.last(); + reuse_token.copied() + } + /** Push a reuse token for a layout on the stack for said layout. */ - fn push_reuse_token(&mut self, layout: &InLayout<'a>, token: ResetToken) { + fn push_reuse_token(&mut self, layout: &InLayout<'a>, token: ReuseToken) { match self.reuse_tokens.get_mut(layout) { Some(reuse_tokens) => { // If the layout is already in the map, push the token on the stack. @@ -362,13 +554,14 @@ impl<'a> ReuseEnvironment<'a> { Add the layout of a symbol. */ fn add_symbol_layout(&mut self, symbol: Symbol, layout: &'a InLayout<'a>) { - self.symbol_layouts.insert(symbol, layout); + self.symbol_layouts + .insert(symbol, LayoutOption::Layout(layout)); } /** Retrieve the layout of a symbol. */ - fn get_symbol_layout(&self, symbol: Symbol) -> &'a InLayout<'a> { + fn get_symbol_layout(&self, symbol: Symbol) -> &LayoutOption<'a> { self.symbol_layouts.get(&symbol).expect("Expected symbol to have a layout. It should have been inserted in the environment already.") } } @@ -407,7 +600,26 @@ fn can_reuse_layout_tag<'a, 'i>( } } }, - // TODO why are strings/lists not reusable? + // Strings literals are constants. + // Arrays are probably given to functions and reused there. Little use to reuse them here. _ => Reuse::Nonreusable, } } + +/** +Drop the reuse tokens that are not used anymore. +Usefull when reuse tokens are used in a branch, and thus should be created. +But not in all branches, and thus should be dropped in those branches. +*/ +fn drop_unused_reuse_tokens<'a>( + arena: &'a Bump, + unused_tokens: impl Iterator, + continuation: &'a Stmt<'a>, +) -> &'a Stmt<'a> { + unused_tokens.fold(continuation, |continuation, reuse_token| { + arena.alloc(Stmt::Refcounting( + ModifyRc::Dec(reuse_token.symbol), + continuation, + )) + }) +} diff --git a/crates/compiler/mono/src/perceus.rs b/crates/compiler/mono/src/perceus.rs index 480ef0f77a..d6e07d28b3 100644 --- a/crates/compiler/mono/src/perceus.rs +++ b/crates/compiler/mono/src/perceus.rs @@ -31,7 +31,6 @@ pub fn insert_refcount_operations<'a, 'i>( // Create a VariableRcTypesEnv for the procedures as they get referenced but should be marked as non reference counted. let mut variable_rc_types_env = VariableRcTypesEnv::from_layout_interner(layout_interner); variable_rc_types_env.insert_proc_symbols(procedures.keys().map(|(symbol, _layout)| *symbol)); - for (_, proc) in procedures.iter_mut() { // Clone the variable_rc_types_env and insert the variables in the current procedure. // As the variables should be limited in scope for the current proc. @@ -297,18 +296,15 @@ impl VariableUsage { variable_rc_types, owned_arguments.iter().map(|(symbol, _)| symbol).copied(), ), - borrowed: borrowed_arguments - .iter() - .map(|(symbol, _)| symbol) - .copied() - .collect(), + borrowed: Self::borrowed_usages( + variable_rc_types, + borrowed_arguments.iter().map(|(symbol, _)| symbol).copied(), + ), }; } CallType::HigherOrder(HigherOrderLowLevel { op: operator, - /// TODO I _think_ we can get rid of this, perhaps only keeping track of - /// the layout of the closure argument, if any closure_env_layout, /// update mode of the higher order lowlevel itself @@ -319,8 +315,12 @@ impl VariableUsage { // Functions always take their arguments as owned. // (Except lowlevels, but those are wrapped in functions that take their arguments as owned and perform rc.) + // TODO do something with these variables. I think there should be only one argument for the closure. let closure_arguments = &arguments[operator.function_index()..]; + // This should always be true, not sure where this could be set to false. + debug_assert!(passed_function.owns_captured_environment); + match operator { crate::low_level::HigherOrder::ListMap { xs } => VariableUsage { owned: Self::owned_usages(variable_rc_types, [*xs].into_iter()), @@ -351,7 +351,9 @@ impl VariableUsage { // But functions assume that they are called with owned arguments, this creates a problem. // We need to treat them as owned by incrementing the reference count before calling the function, { + // TODO probably update the list sort implementation to assume the functions take their argument as owned. todo!() + // TODO sort will perform sort in place (if unique), take this into account. } } } @@ -373,11 +375,7 @@ impl VariableUsage { // VariableUsage::Borrowed(*structure) VariableUsage { owned: MutMap::default(), - borrowed: { - let mut set = MutSet::with_capacity_and_hasher(1, default_hasher()); - set.insert(*structure); - set - }, + borrowed: Self::borrowed_usages(variable_rc_types, iter::once(*structure)), } } Expr::Array { @@ -433,6 +431,30 @@ impl VariableUsage { } variable_usage } + + /** + Filter the given symbols to only contain reference counted symbols. + */ + fn borrowed_usages( + variable_rc_types: &VariableRcTypes, + symbols: impl IntoIterator, + ) -> MutSet { + symbols + .into_iter() + .filter_map(|symbol| { + match { + variable_rc_types + .get(&symbol) + .expect("Expected variable to be in the map") + } { + // If the variable is reference counted, we need to increment the usage count. + VarRcType::ReferenceCounted => Some(symbol), + // If the variable is not reference counted, we don't need to do anything. + VarRcType::NotReferenceCounted => None, + } + }) + .collect() + } } /** @@ -1134,7 +1156,7 @@ fn insert_dec_stmt<'a, 's>( /** Taken from the original inc_dec borrow implementation. */ -pub fn lowlevel_borrow_signature<'a>(arena: &'a Bump, op: &LowLevel) -> &'a [Ownership] { +fn lowlevel_borrow_signature<'a>(arena: &'a Bump, op: &LowLevel) -> &'a [Ownership] { use LowLevel::*; let irrelevant = Ownership::Owned;