inserted reset/reuse

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.
+            let new_expr = match expr {
-            // TODO deal with fresh allocations w/ reuse.
+                Expr::Tag {
-            if let Expr::Tag {
+                    tag_layout,
-                tag_layout,
+                    tag_id,
-                tag_id,
+                    arguments,
-                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.
-                // The value of the tag is currently only used in the case of nullable recursive unions.
+                    environment.add_layout_tag(layout, *tag_id);
-                // 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::<std::vec::Vec<_>>();
 
@@ -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.
+            // First we determine the minimum of reuse tokens available for each layout.
-            // Make sure the current env is updated to have the correct reuse tokens consumed.
+            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
+                    for (_, _, _, branch_env) in new_branches.iter() {
-            todo!()
+                        branch_environments.push(branch_env);
-        }
+                    }
-        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)
 
-                    // Get the layout of the symbol from where it is defined.
+                    branch_environments.push(&new_default_branch.2);
-                    let layout = environment.get_symbol_layout(*symbol);
 
-                    // If the layout is reusable, add a reuse token for the layout to the environment.
+                    branch_environments
-                    if matches!(
+                };
-                        can_reuse_layout_tag(layout_interner, &environment, &layout),
+
-                        Reuse::Reusable
+                let layout_min_reuse_tokens =
-                    ) {
+                    environment.reuse_tokens.keys().copied().map(|layout| {
-                        environment.push_reuse_token(&layout, *symbol);
+                        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::<MutMap<_, _>>()
+            };
+
+            // 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) =>
+        Stmt::Ret(symbol) => {
-        // The return statement just doesn't consume any tokens. Dropping these tokens will be handled before.
+            // 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<InLayout<'a>, Vec<ResetToken>>;
+type ReuseTokens<'a> = MutMap<InLayout<'a>, std::vec::Vec<ReuseToken>>;
 
 /**
 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<Symbol, LayoutOption<'a>>;
+
+#[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, &'a InLayout<'a>>,
+    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<FooBar> {
+    fn get_layout_tag(&self, layout: &InLayout<'a>) -> Option<Tag> {
         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<ResetToken> {
+    fn pop_reuse_token(&mut self, layout: &InLayout<'a>) -> Option<ReuseToken> {
         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<ReuseToken> {
+        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<Item = ReuseToken>,
+    continuation: &'a Stmt<'a>,
+) -> &'a Stmt<'a> {
+    unused_tokens.fold(continuation, |continuation, reuse_token| {
+        arena.alloc(Stmt::Refcounting(
+            ModifyRc::Dec(reuse_token.symbol),
+            continuation,
+        ))
+    })
+}

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
+                        borrowed: Self::borrowed_usages(
-                            .iter()
+                            variable_rc_types,
-                            .map(|(symbol, _)| symbol)
+                            borrowed_arguments.iter().map(|(symbol, _)| symbol).copied(),
-                            .copied()
+                        ),
-                            .collect(),
                     };
                 }
                 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: {
+                    borrowed: Self::borrowed_usages(variable_rc_types, iter::once(*structure)),
-                        let mut set = MutSet::with_capacity_and_hasher(1, default_hasher());
-                        set.insert(*structure);
-                        set
-                    },
                 }
             }
             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<Item = Symbol>,
+    ) -> MutSet<Symbol> {
+        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;