Comments and cleanup

crates/compiler/mono/src/perceus.rs

@@ -7,10 +7,13 @@ use roc_module::symbol::{IdentIds, ModuleId, Symbol};
 
 use crate::{
     borrow::Ownership,
-    ir::{BranchInfo, Call, CallType, Expr, ModifyRc, Proc, ProcLayout, Stmt, UpdateModeIds},
+    ir::{BranchInfo, Expr, ListLiteralElement, ModifyRc, Proc, ProcLayout, Stmt, UpdateModeIds},
     layout::{InLayout, LayoutInterner, STLayoutInterner},
 };
 
+/**
+Insert the reference count operations for procedures.
+*/
 pub fn insert_refcount_operations<'a>(
     arena: &'a Bump,
     layout_interner: &STLayoutInterner,
@@ -20,14 +23,13 @@ pub fn insert_refcount_operations<'a>(
     procedures: &mut HashMap<(Symbol, ProcLayout), Proc<'a>, BuildHasherDefault<WyHash>>,
 ) -> () {
     // Create a VariableRcTypesEnv for the procedures as they get referenced but should be marked as non reference counted.
-    let proc_rc_type = procedures.keys().map(|(symbol, _layout)| *symbol);
+    let mut variable_rc_types_env = VariableRcTypesEnv::from_layout_interner(layout_interner);
-    let variable_rc_types_env =
+    variable_rc_types_env.insert_proc_symbols(procedures.keys().map(|(symbol, _layout)| *symbol));
-        VariableRcTypesEnv::from_proc_symbols(layout_interner, proc_rc_type);
 
     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.
-        let variable_rc_types_env = &mut variable_rc_types_env.clone();
+        let mut variable_rc_types_env = variable_rc_types_env.clone();
         variable_rc_types_env.insert_variables_rc_type_proc(&proc);
 
         let mut initial_environment = Environment {
@@ -56,6 +58,9 @@ type VariableRcTypes = MutMap<Symbol, VarRcType>;
 
 type FreeRcVariables = MutSet<Symbol>;
 
+/**
+Environment to keep track which of the variables should be reference counted and which ones should not.
+ */
 struct VariableRcTypesEnv<'a> {
     // A map keeping track of which variables are reference counted and which are not.
     variables_rc_type: VariableRcTypes,
@@ -65,41 +70,66 @@ struct VariableRcTypesEnv<'a> {
 
 // TODO what would be a good way to structure a similar pattern? creating env, evaluating multiple different objects and returning an element from the env.
 impl<'a> VariableRcTypesEnv<'a> {
-    fn from_proc_symbols(
+    /**
-        layout_interner: &'a STLayoutInterner,
+    Create a new VariableRcTypesEnv from a layout interner.
-        proc_symbols: impl Iterator<Item = Symbol>,
+    */
-    ) -> VariableRcTypesEnv<'a> {
+    fn from_layout_interner(layout_interner: &'a STLayoutInterner) -> VariableRcTypesEnv<'a> {
         VariableRcTypesEnv {
-            variables_rc_type: proc_symbols
+            variables_rc_type: VariableRcTypes::default(),
-                .map(|symbol| (symbol, VarRcType::NotReferenceCounted))
-                .collect(),
             layout_interner,
         }
     }
+    /**
+    Insert the reference count type of top level functions.
+    As functions are not reference counted, they can be marked as such.
+    */
+    fn insert_proc_symbols(
+        self: &mut VariableRcTypesEnv<'a>,
+        proc_symbols: impl Iterator<Item = Symbol>,
+    ) {
+        for proc_symbol in proc_symbols {
+            self.variables_rc_type
+                .insert(proc_symbol, VarRcType::NotReferenceCounted);
+        }
+    }
 
-    // Get the reference count types of all variables in a procedure.
+    /**
+    Insert the reference count types of all variables in a procedure.
+    */
     fn insert_variables_rc_type_proc(self: &mut VariableRcTypesEnv<'a>, proc: &Proc<'a>) {
+        // First collect the argument types.
         for (layout, symbol) in proc.args.iter() {
             self.insert_symbol_layout_rc_type(symbol, layout);
         }
 
+        // Then collect the types of the variables in the body.
         self.insert_variables_rc_type_stmt(&proc.body);
     }
 
+    /**
+    Insert the reference count types of all variables in a statement.
+    */
     fn insert_variables_rc_type_stmt(self: &mut VariableRcTypesEnv<'a>, stmt: &Stmt<'a>) {
         match stmt {
-            Stmt::Let(binding, _expr, layout, stmt) => {
+            Stmt::Let(
+                binding,
+                // Expressions can be omitted, as they won't create new variables.
+                _expr,
+                layout,
+                continuation,
+            ) => {
                 self.insert_symbol_layout_rc_type(&binding, layout);
-                self.insert_variables_rc_type_stmt(stmt);
+                self.insert_variables_rc_type_stmt(continuation);
             }
             Stmt::Switch {
-                // TODO what is this symbol for and should it be handled?
+                // The switch condition is an integer and thus not reference counted.
-                cond_symbol,
+                cond_symbol: _,
-                cond_layout,
+                cond_layout: _,
                 branches,
                 default_branch,
                 ret_layout: _,
             } => {
+                // Collect the types of the variables in all the branches, including the default one.
                 for (info, stmt) in branches
                     .iter()
                     .map(|(_branch, info, stmt)| (info, stmt))
@@ -119,8 +149,12 @@ impl<'a> VariableRcTypesEnv<'a> {
                     self.insert_variables_rc_type_stmt(stmt);
                 }
             }
-            Stmt::Ret(_) => {}
+            Stmt::Ret(_symbol) => {
-            Stmt::Refcounting(_, _) => {}
+                // The return does not introduce new variables.
+            }
+            Stmt::Refcounting(_, _) => unreachable!(
+                "Refcounting operations should not be present in the AST at this point."
+            ),
             Stmt::Expect {
                 condition,
                 region,
@@ -128,6 +162,7 @@ impl<'a> VariableRcTypesEnv<'a> {
                 variables,
                 remainder,
             } => {
+                //TODO deal with other variables in the expect.
                 self.insert_variables_rc_type_stmt(remainder);
             }
             Stmt::ExpectFx {
@@ -137,6 +172,7 @@ impl<'a> VariableRcTypesEnv<'a> {
                 variables,
                 remainder,
             } => {
+                //TODO deal with other variables in the expect.
                 self.insert_variables_rc_type_stmt(remainder);
             }
             Stmt::Dbg {
@@ -144,19 +180,35 @@ impl<'a> VariableRcTypesEnv<'a> {
                 variable,
                 remainder,
             } => {
+                //TODO deal with other variables in the expect.
                 self.insert_variables_rc_type_stmt(remainder);
             }
             Stmt::Join {
-                id,
+                id: _,
                 parameters,
                 body,
-                remainder,
+                remainder: continuation,
-            } => todo!(),
+            } => {
-            Stmt::Jump(_, _) => todo!(),
+                // TODO do we need to do anything with the parameter ownership parameter?
-            Stmt::Crash(_, _) => todo!(),
+                for parameter in parameters.iter() {
+                    self.insert_symbol_layout_rc_type(&parameter.symbol, &parameter.layout);
+                }
+
+                self.insert_variables_rc_type_stmt(body);
+                self.insert_variables_rc_type_stmt(continuation);
+            }
+            Stmt::Jump(_, _) => {
+                // A join point does not introduce new variables.
+            }
+            Stmt::Crash(_, _) => {
+                // A crash does not introduce new variables.
+            }
         }
     }
 
+    /*
+    Insert the reference count type of a symbol given its layout.
+    */
     fn insert_symbol_layout_rc_type(
         self: &mut VariableRcTypesEnv<'a>,
         symbol: &Symbol,
@@ -201,6 +253,9 @@ fn combine_variable_usage(
     }
 }
 
+/**
+A struct to determine the usage of variables in an expression.
+ */
 struct VariableUsageEnv<'a> {
     variable_usage: VariableUsage,
 
@@ -208,10 +263,13 @@ struct VariableUsageEnv<'a> {
 }
 
 impl<'a> VariableUsageEnv<'a> {
+    /**
+    Retrieve how much a variable is used in an expression.
+    */
     fn get_reference_counted_variable_usage_expr(
         variable_rc_types: &VariableRcTypes,
         expr: &Expr<'a>,
-    ) -> HashMap<Symbol, u64, BuildHasherDefault<WyHash>> {
+    ) -> VariableUsage {
         let mut usage_env = VariableUsageEnv {
             variable_usage: VariableUsage::default(),
             variable_rc_types,
@@ -228,8 +286,8 @@ impl<'a> VariableUsageEnv<'a> {
                 usage_env.insert_variable_usages(arguments.iter().copied());
             }
             Expr::StructAtIndex {
-                index,
+                index: _,
-                field_layouts,
+                field_layouts: _,
                 structure,
             } => {
                 usage_env.insert_variable_usage(*structure);
@@ -246,30 +304,41 @@ impl<'a> VariableUsageEnv<'a> {
                 union_layout,
                 index,
             } => todo!(),
-            Expr::Array { elem_layout, elems } => todo!(),
+            Expr::Array {
-            Expr::EmptyArray => todo!(),
+                elem_layout: _,
+                elems,
+            } => {
+                // For an array creation, we insert all the used elements.
+                usage_env.insert_variable_usages(elems.iter().filter_map(
+                    |element| match element {
+                        // Literal elements are not reference counted.
+                        ListLiteralElement::Literal(_) => None,
+                        // Symbol elements are reference counted.
+                        ListLiteralElement::Symbol(symbol) => Some(*symbol),
+                    },
+                ));
+            }
+            Expr::EmptyArray => {
+                // Empty arrays have no reference counted elements.
+            }
             Expr::ExprBox { symbol } => todo!(),
             Expr::ExprUnbox { symbol } => todo!(),
-            Expr::Reuse {
+            Expr::Reuse { .. } | Expr::Reset { .. } => {
-                symbol,
+                unreachable!("Reset and reuse should not exist at this point")
-                update_tag_id,
+            }
-                update_mode,
-                tag_layout,
-                tag_id,
-                arguments,
-            } => todo!(),
-            Expr::Reset {
-                symbol,
-                update_mode,
-            } => todo!(),
             Expr::RuntimeErrorFunction(_) => todo!(),
         }
 
         usage_env.variable_usage
     }
 
+    /**
+    Insert the usage of a symbol into the env.
+    If the symbol is not reference counted, it will be ignored.
+    */
     fn insert_variable_usage(self: &mut VariableUsageEnv<'a>, symbol: Symbol) {
         match {
+            // TODO convert this to a oneliner once panics do not occur anymore.
             let this = self.variable_rc_types.get(&symbol);
             match this {
                 Some(val) => val,
@@ -288,6 +357,9 @@ impl<'a> VariableUsageEnv<'a> {
         }
     }
 
+    /**
+    Call insert_variable_usage for multiple symbols.
+    */
     fn insert_variable_usages(
         self: &mut VariableUsageEnv<'a>,
         symbols: impl Iterator<Item = Symbol>,
@@ -296,6 +368,10 @@ impl<'a> VariableUsageEnv<'a> {
     }
 }
 
+/**
+The environment for the reference counting pass.
+Contains the variable rc types and the ownership.
+*/
 struct Environment<'a> {
     // Keep track which variables are reference counted and which are not.
     variables_rc_types: &'a VariableRcTypes,
@@ -313,14 +389,19 @@ impl<'a> Clone for Environment<'a> {
 }
 
 impl<'a> Environment<'a> {
+    /**
+    Retrieve the rc type of a variable.
+    */
     fn get_variable_rc_type(self: &mut Environment<'a>, variable: &Symbol) -> &VarRcType {
         self.variables_rc_types
             .get(variable)
             .expect("variable should have rc type")
     }
 
-    // Retrieve whether the variable is owned or borrowed.
+    /*
-    // If it was owned, set it to borrowed.
+    Retrieve whether the variable is owned or borrowed.
+    If it was owned, set it to borrowed (as it is consumed/the variable can be used as owned only once without incrementing).
+    */
     fn consume_variable(self: &mut Environment<'a>, variable: &Symbol) -> Ownership {
         // This function should only be called on reference counted variables.
         debug_assert!(matches!(
@@ -337,7 +418,9 @@ impl<'a> Environment<'a> {
         }
     }
 
-    // Add a variable to the environment.
+    /*
+    Add a variable to the environment if it is reference counted.
+    */
     fn add_variable(self: &mut Environment<'a>, variable: &Symbol) {
         match self.get_variable_rc_type(variable) {
             VarRcType::ReferenceCounted => {
@@ -349,8 +432,10 @@ impl<'a> Environment<'a> {
         }
     }
 
-    // Add a variable to the environment during a function call.
+    /*
-    // Remove the variable afterwards to prevent it from being used outside the function call.
+    Add a variable to the environment during a function call.
+    Remove the variable afterwards to prevent it from being used outside the function call.
+    */
     fn with_variable<F, R>(self: &mut Environment<'a>, variable: &Symbol, callback: F) -> R
     where
         F: Fn(&mut Environment) -> R,
@@ -367,8 +452,9 @@ impl<'a> Environment<'a> {
     }
 }
 
-// TODO take into account what values should and should not be reference counted.
+/**
-// e.g. stack allocated values should not be reference counted.
+Given an environment, insert the reference counting operations for a statement.
+*/
 fn insert_refcount_operations_stmt<'a>(
     arena: &'a Bump,
     environment: &mut Environment,
@@ -439,6 +525,7 @@ fn insert_refcount_operations_stmt<'a>(
                 (info.clone(), new_branch, free_variables_branch)
             };
 
+            // Combine the free variables of the all branches.
             let free_rc_vars = {
                 let mut free_rc_vars = FreeRcVariables::default();
 
@@ -455,6 +542,9 @@ fn insert_refcount_operations_stmt<'a>(
                 free_rc_vars
             };
 
+            // If we compare the free variables from a branch with the free variables from all branches.
+            // We can determine which variables are not used in that branch.
+            // And as such, we can drop them before continuing.
             let newer_branches = Vec::from_iter_in(
                 new_branches
                     .into_iter()
@@ -486,8 +576,13 @@ fn insert_refcount_operations_stmt<'a>(
             (new_switch, free_rc_vars)
         }
         Stmt::Ret(s) => {
-            return (arena.alloc(Stmt::Ret(*s)), FreeRcVariables::default());
+            // TODO use with_capacity if possible, or use an initializer to oneline this.
-            todo!()
+            let mut free_variables = FreeRcVariables::with_capacity_and_hasher(
+                1,
+                roc_collections::all::BuildHasher::default(),
+            );
+            free_variables.insert(*s);
+            return (arena.alloc(Stmt::Ret(*s)), free_variables);
         }
         Stmt::Refcounting(_, _) => unreachable!("refcounting should not be in the AST already"),
         Stmt::Expect {
@@ -520,6 +615,9 @@ fn insert_refcount_operations_stmt<'a>(
     }
 }
 
+/**
+ * Insert increment statements for the given symbols compensating for the ownership.
+ */
 fn insert_inc_stmts<'a>(
     arena: &'a Bump,
     environment: &mut Environment,
@@ -533,6 +631,9 @@ fn insert_inc_stmts<'a>(
         })
 }
 
+/**
+Insert an increment statement for the given symbol compensating for the ownership.
+*/
 fn insert_inc_stmt<'a>(
     arena: &'a Bump,
     environment: &mut Environment,
@@ -556,6 +657,9 @@ fn insert_inc_stmt<'a>(
     }
 }
 
+/**
+Insert decrement statements for the given symbols if they are owned.
+*/
 fn insert_dec_stmts<'a>(
     arena: &'a Bump,
     environment: &mut Environment,
@@ -567,6 +671,9 @@ fn insert_dec_stmts<'a>(
     })
 }
 
+/**
+Insert a decrement statement for the given symbol if it is owned.
+*/
 fn insert_dec_stmt<'a>(
     arena: &'a Bump,
     environment: &mut Environment,