Start new free var inference and joint points

crates/compiler/mono/src/perceus.rs

@@ -7,16 +7,19 @@ use roc_module::symbol::{IdentIds, ModuleId, Symbol};
 
 use crate::{
     borrow::Ownership,
-    ir::{BranchInfo, Expr, ListLiteralElement, ModifyRc, Proc, ProcLayout, Stmt, UpdateModeIds},
+    ir::{
+        BranchInfo, Call, Expr, JoinPointId, ListLiteralElement, ModifyRc, Param, Proc, ProcLayout,
+        Stmt, UpdateModeIds,
+    },
     layout::{InLayout, LayoutInterner, STLayoutInterner},
 };
 
 /**
 Insert the reference count operations for procedures.
 */
-pub fn insert_refcount_operations<'a>(
+pub fn insert_refcount_operations<'a, 'i>(
     arena: &'a Bump,
-    layout_interner: &STLayoutInterner,
+    layout_interner: &'i STLayoutInterner<'a>,
     home: ModuleId,
     ident_ids: &mut IdentIds,
     update_mode_ids: &mut UpdateModeIds,
@@ -29,25 +32,16 @@ pub fn insert_refcount_operations<'a>(
     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 mut variable_rc_types_env = variable_rc_types_env.clone();
+        let variable_rc_types_env = variable_rc_types_env.clone();
-        variable_rc_types_env.insert_variables_rc_type_proc(&proc);
 
-        let mut initial_environment = Environment {
+        insert_refcount_operations_proc(arena, variable_rc_types_env, proc);
-            variables_rc_types: &variable_rc_types_env.variables_rc_type,
-            variables_ownership: MutMap::default(),
-        };
-
-        // Add all arguments to the environment (if they are reference counted)
-        for (_layout, symbol) in proc.args.iter() {
-            initial_environment.add_variable(symbol);
-        }
-
-        let (new_body, _) =
-            insert_refcount_operations_stmt(arena, &mut initial_environment, &proc.body);
-        proc.body = new_body.clone();
     }
 }
 
+/**
+Enum indicating whether a variable (symbol) should be reference counted or not.
+This includes layouts that themselves can be stack allocated but that contain a heap allocated item.
+*/
 #[derive(Clone)]
 enum VarRcType {
     ReferenceCounted,
@@ -56,24 +50,24 @@ enum VarRcType {
 
 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> {
+struct VariableRcTypesEnv<'a, 'i> {
     // A map keeping track of which variables are reference counted and which are not.
     variables_rc_type: VariableRcTypes,
 
-    layout_interner: &'a STLayoutInterner<'a>,
+    layout_interner: &'i STLayoutInterner<'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> {
+impl<'a, 'i> VariableRcTypesEnv<'a, 'i> {
     /**
     Create a new VariableRcTypesEnv from a layout interner.
     */
-    fn from_layout_interner(layout_interner: &'a STLayoutInterner) -> VariableRcTypesEnv<'a> {
+    fn from_layout_interner(
+        layout_interner: &'i STLayoutInterner<'a>,
+    ) -> VariableRcTypesEnv<'a, 'i> {
         VariableRcTypesEnv {
             variables_rc_type: VariableRcTypes::default(),
             layout_interner,
@@ -84,7 +78,7 @@ impl<'a> VariableRcTypesEnv<'a> {
     As functions are not reference counted, they can be marked as such.
     */
     fn insert_proc_symbols(
-        self: &mut VariableRcTypesEnv<'a>,
+        self: &mut VariableRcTypesEnv<'a, 'i>,
         proc_symbols: impl Iterator<Item = Symbol>,
     ) {
         for proc_symbol in proc_symbols {
@@ -96,7 +90,7 @@ impl<'a> VariableRcTypesEnv<'a> {
     /**
     Insert the reference count types of all variables in a procedure.
     */
-    fn insert_variables_rc_type_proc(self: &mut VariableRcTypesEnv<'a>, proc: &Proc<'a>) {
+    fn insert_variables_rc_type_proc(self: &mut VariableRcTypesEnv<'a, 'i>, proc: &Proc<'a>) {
         // First collect the argument types.
         for (layout, symbol) in proc.args.iter() {
             self.insert_symbol_layout_rc_type(symbol, layout);
@@ -109,7 +103,7 @@ impl<'a> VariableRcTypesEnv<'a> {
     /**
     Insert the reference count types of all variables in a statement.
     */
-    fn insert_variables_rc_type_stmt(self: &mut VariableRcTypesEnv<'a>, stmt: &Stmt<'a>) {
+    fn insert_variables_rc_type_stmt(self: &mut VariableRcTypesEnv<'a, 'i>, stmt: &Stmt<'a>) {
         match stmt {
             Stmt::Let(
                 binding,
@@ -210,7 +204,7 @@ impl<'a> VariableRcTypesEnv<'a> {
     Insert the reference count type of a symbol given its layout.
     */
     fn insert_symbol_layout_rc_type(
-        self: &mut VariableRcTypesEnv<'a>,
+        self: &mut VariableRcTypesEnv<'a, 'i>,
         symbol: &Symbol,
         layout: &InLayout,
     ) {
@@ -228,7 +222,7 @@ impl<'a> VariableRcTypesEnv<'a> {
     }
 }
 
-impl Clone for VariableRcTypesEnv<'_> {
+impl Clone for VariableRcTypesEnv<'_, '_> {
     fn clone(&self) -> Self {
         VariableRcTypesEnv {
             variables_rc_type: self.variables_rc_type.clone(),
@@ -239,151 +233,158 @@ impl Clone for VariableRcTypesEnv<'_> {
 
 type VariablesOwnership = MutMap<Symbol, Ownership>;
 
-// A map keeping track of how many times a variable is used.
+/**
-type VariableUsage = MutMap<Symbol, u64>;
+A map keeping track of how many times a variable is used as owned.
-
+*/
-// Combine variable usage by summing the usage of each variable.
+type OwnedUsage = MutMap<Symbol, u64>;
-fn combine_variable_usage(
-    variable_usage: &mut VariableUsage,
-    other_variable_usage: &VariableUsage,
-) {
-    for (symbol, other_usage) in other_variable_usage.iter() {
-        match variable_usage.get(symbol) {
-            Some(old_usage) => variable_usage.insert(*symbol, old_usage + other_usage),
-            None => variable_usage.insert(*symbol, *other_usage),
-        };
-    }
-}
 
 /**
-A struct to determine the usage of variables in an expression.
+Structure to keep track of the borrowed variable usage of an expression.
- */
+For now a single symbol as struct/union indexing always happens on just a single item.
-struct VariableUsageEnv<'a> {
+*/
-    variable_usage: VariableUsage,
+type BorrowedUsage = Symbol;
 
-    variable_rc_types: &'a VariableRcTypes,
+/**
+Enum to keep track of the variable usage of an expression.
+*/
+enum VariableUsage {
+    None,
+    Owned(OwnedUsage),
+    Borrowed(BorrowedUsage),
 }
 
-impl<'a> VariableUsageEnv<'a> {
+impl VariableUsage {
     /**
     Retrieve how much a variable is used in an expression.
     */
-    fn get_reference_counted_variable_usage_expr(
+    fn get_reference_counted_variable_usage_expr<'a>(
         variable_rc_types: &VariableRcTypes,
         expr: &Expr<'a>,
     ) -> VariableUsage {
-        let mut usage_env = VariableUsageEnv {
-            variable_usage: VariableUsage::default(),
-            variable_rc_types,
-        };
-
         match expr {
             Expr::Literal(_) => {
                 // Literals are not reference counted.
+                VariableUsage::None
             }
-            Expr::Call(call) => {
+
-                usage_env.insert_variable_usages(call.arguments.iter().copied());
+            Expr::Call(Call { arguments, .. })
-            }
+            | Expr::Tag { arguments, .. }
-            Expr::Tag { arguments, .. } | Expr::Struct(arguments) => {
+            | Expr::Struct(arguments) => {
-                usage_env.insert_variable_usages(arguments.iter().copied());
+                Self::owned_usages(variable_rc_types, arguments.iter().copied())
             }
 
             Expr::GetTagId { structure, .. }
             | Expr::StructAtIndex { structure, .. }
             | Expr::UnionAtIndex { structure, .. } => {
                 // All structures are alive at this point and don't have to be copied in order to take an index out.
-                // TODO perhaps inc if the index is a pointer to a heap value, assuming non pointers get copied.
+                // But we do want to make sure to decrement this item if it is the last reference.
+                VariableUsage::Borrowed(*structure)
             }
             Expr::Array {
                 elem_layout: _,
                 elems,
             } => {
                 // For an array creation, we insert all the used elements.
-                usage_env.insert_variable_usages(elems.iter().filter_map(
+                Self::owned_usages(
-                    |element| match element {
+                    variable_rc_types,
+                    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.
+                VariableUsage::None
             }
             Expr::ExprBox { symbol } | Expr::ExprUnbox { symbol } => {
-                usage_env.insert_variable_usage(*symbol);
+                Self::owned_usages(variable_rc_types, iter::once(*symbol))
             }
             Expr::Reuse { .. } | Expr::Reset { .. } => {
                 unreachable!("Reset and reuse should not exist at this point")
             }
             Expr::RuntimeErrorFunction(_) => todo!(),
         }
-
-        usage_env.variable_usage
     }
 
     /**
-    Insert the usage of a symbol into the env.
+    Return owned usages.
-    If the symbol is not reference counted, it will be ignored.
+    Collect the usage of all the reference counted symbols in the iterator and return as a map.
     */
-    fn insert_variable_usage(self: &mut VariableUsageEnv<'a>, symbol: Symbol) {
+    fn owned_usages(
-        match {
+        variable_rc_types: &VariableRcTypes,
-            // 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,
-                None => panic!("Expected variable to be in the map"),
-            }
-        } {
-            // If the variable is reference counted, we need to increment the usage count.
-            VarRcType::ReferenceCounted => {
-                match self.variable_usage.get(&symbol) {
-                    Some(count) => self.variable_usage.insert(symbol, count + 1),
-                    None => self.variable_usage.insert(symbol, 1),
-                };
-            }
-            // If the variable is not reference counted, we don't need to do anything.
-            VarRcType::NotReferenceCounted => return,
-        }
-    }
-
-    /**
-    Call insert_variable_usage for multiple symbols.
-    */
-    fn insert_variable_usages(
-        self: &mut VariableUsageEnv<'a>,
         symbols: impl Iterator<Item = Symbol>,
-    ) {
+    ) -> VariableUsage {
-        symbols.for_each(|symbol| self.insert_variable_usage(symbol));
+        // A groupby or something similar would be nice here.
+        let mut variable_usage = OwnedUsage::default();
+        symbols.for_each(|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 => {
+                    match variable_usage.get(&symbol) {
+                        Some(count) => variable_usage.insert(symbol, count + 1),
+                        None => variable_usage.insert(symbol, 1),
+                    };
+                }
+                // If the variable is not reference counted, we don't need to do anything.
+                VarRcType::NotReferenceCounted => return,
+            }
+        });
+        VariableUsage::Owned(variable_usage)
     }
 }
 
+/**
+Enum to indicate whether or not a join point consumed a parameter.
+*/
+enum Consumption {
+    Consumed,
+    Unconsumed,
+}
+
+/**
+Struct containing data about the variable consumption of a join point.
+Contains data about consumed closure values and the consumption of the parameters.
+*/
+#[derive(Clone)]
+struct JoinPointConsumption<'a> {
+    closure: MutSet<Symbol>,
+    parameters: &'a [Consumption],
+}
+
 /**
 The environment for the reference counting pass.
 Contains the variable rc types and the ownership.
 */
-struct Environment<'a> {
+struct Environment<'v, 'a> {
     // Keep track which variables are reference counted and which are not.
-    variables_rc_types: &'a VariableRcTypes,
+    variables_rc_types: &'v VariableRcTypes,
     // The Koka implementation assumes everything that is not owned to be borrowed.
     variables_ownership: VariablesOwnership,
+    jointpoint_closures: MutMap<JoinPointId, JoinPointConsumption<'a>>,
 }
 
-impl<'a> Clone for Environment<'a> {
+impl<'v, 'a> Clone for Environment<'v, 'a> {
     fn clone(&self) -> Self {
         Environment {
             variables_rc_types: self.variables_rc_types,
             variables_ownership: self.variables_ownership.clone(),
+            jointpoint_closures: self.jointpoint_closures.clone(),
         }
     }
 }
 
-impl<'a> Environment<'a> {
+impl<'v, 'a> Environment<'v, 'a> {
     /**
     Retrieve the rc type of a variable.
     */
-    fn get_variable_rc_type(self: &mut Environment<'a>, variable: &Symbol) -> &VarRcType {
+    fn get_variable_rc_type(self: &mut Environment<'v, 'a>, variable: &Symbol) -> &VarRcType {
         self.variables_rc_types
             .get(variable)
             .expect("variable should have rc type")
@@ -392,30 +393,50 @@ impl<'a> Environment<'a> {
     /*
     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).
+    If the variable is not reference counted, do nothing and return None.
     */
-    fn consume_variable(self: &mut Environment<'a>, variable: &Symbol) -> Ownership {
+    fn consume_variable(self: &mut Environment<'v, 'a>, variable: &Symbol) -> Option<Ownership> {
-        // This function should only be called on reference counted variables.
+        if !self.variables_ownership.contains_key(variable) {
-        debug_assert!(matches!(
+            return None;
-            self.get_variable_rc_type(variable),
-            VarRcType::ReferenceCounted
-        ));
-        // Consume the variable by setting it to borrowed (if it was owned before), and return the previous ownership.
-        let this = self
-            .variables_ownership
-            .insert(*variable, Ownership::Borrowed);
-        match this {
-            Some(val) => val,
-            None => panic!("Expected variable to be in environment"),
         }
+
+        // Consume the variable and return the previous ownership.
+        Some(self.consume_rc_variable(variable))
     }
 
     /*
+    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_rc_variable(self: &mut Environment<'v, 'a>, variable: &Symbol) -> Ownership {
+        // Consume the variable by setting it to borrowed (if it was owned before), and return the previous ownership.
+        self.variables_ownership
+            .insert(*variable, Ownership::Borrowed)
+            .expect("Expected variable to be in environment")
+    }
+
+    /**
+       Retrieve the ownership of a variable.
+       If the variable is not reference counted, it will None.
+    */
+    fn get_variable_ownership(self: &Environment<'v, 'a>, variable: &Symbol) -> Option<&Ownership> {
+        self.variables_ownership.get(variable)
+    }
+
+    /**
+    Add a variables to the environment if they are reference counted.
+    */
+    fn add_variables(self: &mut Environment<'v, 'a>, variables: impl Iterator<Item = Symbol>) {
+        variables.for_each(|variable| self.add_variable(variable))
+    }
+
+    /**
     Add a variable to the environment if it is reference counted.
     */
-    fn add_variable(self: &mut Environment<'a>, variable: &Symbol) {
+    fn add_variable(self: &mut Environment<'v, 'a>, variable: Symbol) {
-        match self.get_variable_rc_type(variable) {
+        match self.get_variable_rc_type(&variable) {
             VarRcType::ReferenceCounted => {
-                self.variables_ownership.insert(*variable, Ownership::Owned);
+                self.variables_ownership.insert(variable, Ownership::Owned);
             }
             VarRcType::NotReferenceCounted => {
                 // If this variable is not reference counted, we don't need to do anything.
@@ -423,73 +444,163 @@ impl<'a> Environment<'a> {
         }
     }
 
-    /*
+    /**
-    Add a variable to the environment during a function call.
+    Remove variables from the environment.
-    Remove the variable afterwards to prevent it from being used outside the function call.
+    Is used when a variable is no longer in scope (after checking a join point).
+     */
+    fn remove_variables(self: &mut Environment<'v, 'a>, variables: impl Iterator<Item = Symbol>) {
+        variables.for_each(|variable| self.remove_variable(variable))
+    }
+    /**
+    Remove a variable from the environment.
+    Is used when a variable is no longer in scope (before a let binding).
+     */
+    fn remove_variable(self: &mut Environment<'v, 'a>, variable: Symbol) {
+        self.variables_ownership.remove(&variable);
+    }
+
+    /**
+    Add a joinpoint id and the consumed closure to the environment.
+    Used when analyzing a join point. So that a jump can update the environment on call.
     */
-    fn with_variable<F, R>(self: &mut Environment<'a>, variable: &Symbol, callback: F) -> R
+    fn add_joinpoint_consumption(
-    where
+        self: &mut Environment<'v, 'a>,
-        F: Fn(&mut Environment) -> R,
+        joinpoint_id: JoinPointId,
-    {
+        consumption: JoinPointConsumption<'a>,
-        match self.get_variable_rc_type(variable) {
+    ) {
-            VarRcType::ReferenceCounted => {
+        self.jointpoint_closures.insert(joinpoint_id, consumption);
-                self.variables_ownership.insert(*variable, Ownership::Owned);
+    }
-                let result = callback(self);
+
-                self.variables_ownership.remove(&variable);
+    /**
-                result
+    Get the consumed closure from a join point id.
-            }
+    */
-            VarRcType::NotReferenceCounted => callback(self),
+    fn get_joinpoint_consume(
-        }
+        self: &Environment<'v, 'a>,
+        joinpoint_id: JoinPointId,
+    ) -> JoinPointConsumption<'a> {
+        self.jointpoint_closures
+            .get(&joinpoint_id)
+            .expect("Expected closure to be in environment")
+            .clone()
+    }
+
+    /**
+    Remove a joinpoint id and the consumed closure from the environment.
+    Used after analyzing the continuation of a join point.
+    */
+    fn remove_joinpoint_consumption(self: &mut Environment<'v, 'a>, joinpoint_id: JoinPointId) {
+        let closure = self.jointpoint_closures.remove(&joinpoint_id);
+        debug_assert!(
+            matches!(closure, Some(_)),
+            "Expected closure to be in environment"
+        );
     }
 }
 
 /**
-Given an environment, insert the reference counting operations for a statement.
+ Insert the reference counting operations into a statement.
 */
-fn insert_refcount_operations_stmt<'a>(
+fn insert_refcount_operations_proc<'a, 'i>(
     arena: &'a Bump,
-    environment: &mut Environment,
+    mut variable_rc_types_env: VariableRcTypesEnv<'a, 'i>,
+    proc: &mut Proc<'a>,
+) {
+    // 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.
+    variable_rc_types_env.insert_variables_rc_type_proc(&proc);
+
+    let mut environment = Environment {
+        variables_rc_types: &variable_rc_types_env.variables_rc_type,
+        variables_ownership: MutMap::default(),
+        jointpoint_closures: MutMap::default(),
+    };
+
+    // Add all arguments to the environment (if they are reference counted)
+    let mut proc_symbols = proc.args.iter().map(|(_layout, symbol)| symbol);
+    for symbol in proc_symbols.by_ref() {
+        environment.add_variable(*symbol);
+    }
+
+    // Update the body with reference count statements.
+    let new_body = insert_refcount_operations_stmt(arena, &mut environment, &proc.body);
+
+    // Insert decrement statements for unused parameters (which are still marked as owned).
+    let newer_body = consume_and_insert_dec_stmts(arena, &mut environment, proc_symbols, new_body);
+
+    // Assert that just the arguments are in the environment. And (after decrementing the unused ones) that they are all borrowed.
+    debug_assert!(environment
+        .variables_ownership
+        .iter()
+        .all(|(symbol, ownership)| {
+            // All variables should be borrowed.
+            *ownership == Ownership::Borrowed
+                && proc
+                    .args
+                    .iter()
+                    .map(|(_layout, symbol)| symbol)
+                    .any(|s| s == symbol)
+        }));
+
+    proc.body = newer_body.clone();
+}
+
+/**
+Given an environment, insert the reference counting operations for a statement.
+Assuming that a symbol can only be defined once (no binding to the same variable mutliple times).
+*/
+fn insert_refcount_operations_stmt<'v, 'a>(
+    arena: &'a Bump,
+    environment: &mut Environment<'v, 'a>,
     stmt: &Stmt<'a>,
-) -> (&'a Stmt<'a>, FreeRcVariables) {
+) -> &'a Stmt<'a> {
     // TODO: Deal with potentially stack overflowing let chains with an explicit loop.
     // Koka has a list for let bindings.
 
     match &stmt {
         // The expression borrows the values owned (used) by the continuation.
         Stmt::Let(binding, expr, layout, stmt) => {
-            // TODO take into account that the bound variable might not be free in the continuation.
-            // And as such, we can drop the value before continuing.
-
             // INFO The Koka implementation (instead of calculating the owned environment beforehand)
             // First evaluates the continuation with the owned environment, setting the variables to dead (not alive).
             // And in the rest of the code, the dead variables are treated as borrowed (because not alive).
 
-            // First evalute the continuation and let it consume it's free variables.
+            // First evaluate the continuation and let it consume it's free variables.
-            let (new_stmt, mut free_rc_vars_stmt) = environment.with_variable(binding, |env| {
+            environment.add_variable(*binding); // Add the bound variable to the environment. As it can be used in the continuation.
-                insert_refcount_operations_stmt(arena, env, stmt)
+            let mut new_stmt = insert_refcount_operations_stmt(arena, environment, stmt);
-            });
+
+            // If the binding is still owned in the environment, it is not used in the continuation and we can drop it right away.
+            if matches!(
+                environment.get_variable_ownership(binding),
+                Some(Ownership::Owned)
+            ) {
+                new_stmt = insert_dec_stmt(arena, *binding, new_stmt);
+            }
+
+            // And as the variable should not be in scope before this let binding, remove it from the environment.
+            environment.remove_variable(*binding);
 
             // Then evaluate the bound expression. where the free variables from the continuation are borrowed.
-            let variable_usage = VariableUsageEnv::get_reference_counted_variable_usage_expr(
+            let variable_usage = VariableUsage::get_reference_counted_variable_usage_expr(
                 &environment.variables_rc_types,
                 expr,
             );
 
-            let new_let = arena.alloc(Stmt::Let(*binding, expr.clone(), *layout, new_stmt));
+            match variable_usage {
+                VariableUsage::None => {
+                    arena.alloc(Stmt::Let(*binding, expr.clone(), *layout, new_stmt))
+                }
+                VariableUsage::Owned(owned_usage) => {
+                    let new_let = arena.alloc(Stmt::Let(*binding, expr.clone(), *layout, new_stmt));
 
-            // Insert the reference count operations for the variables used in the expression.
+                    // Insert the reference count operations for the owned variables used in the expression.
-            let new_let_with_refcount =
+                    consume_and_insert_inc_stmts(arena, environment, &owned_usage, new_let)
-                insert_inc_stmts(arena, environment, &variable_usage, new_let);
+                }
+                VariableUsage::Borrowed(symbol) => {
+                    let newer_stmt =
+                        consume_and_insert_dec_stmt(arena, environment, &symbol, new_stmt);
 
-            let free_rc_vars = {
+                    arena.alloc(Stmt::Let(*binding, expr.clone(), *layout, newer_stmt))
-                // Remove the bound variable, as it is no longer free.
+                }
-                free_rc_vars_stmt.remove(binding);
+            }
-                // Add the free variables from the expression.
-                free_rc_vars_stmt.extend(variable_usage.keys());
-                free_rc_vars_stmt
-            };
-
-            (new_let_with_refcount, free_rc_vars)
         }
         Stmt::Switch {
             cond_symbol,
@@ -498,50 +609,79 @@ fn insert_refcount_operations_stmt<'a>(
             default_branch,
             ret_layout,
         } => {
-            let new_branches: std::vec::Vec<_> = branches
+            let new_branches = branches
                 .iter()
                 .map(|(label, info, branch)| {
-                    let (new_branch, free_variables_branch) =
+                    let mut branch_env = environment.clone();
-                        insert_refcount_operations_stmt(arena, &mut environment.clone(), branch);
 
-                    (*label, info.clone(), new_branch, free_variables_branch)
+                    let new_branch =
+                        insert_refcount_operations_stmt(arena, &mut branch_env, branch);
+
+                    (*label, info.clone(), new_branch, branch_env)
                 })
-                .collect();
+                .collect::<std::vec::Vec<_>>();
 
             let new_default_branch = {
                 let (info, branch) = default_branch;
-                let (new_branch, free_variables_branch) =
-                    insert_refcount_operations_stmt(arena, &mut environment.clone(), branch);
 
-                (info.clone(), new_branch, free_variables_branch)
+                let mut branch_env = environment.clone();
+                let new_branch = insert_refcount_operations_stmt(arena, &mut branch_env, branch);
+
+                (info.clone(), new_branch, branch_env)
             };
 
-            // Combine the free variables of the all branches.
+            // Determine what variables are consumed in some of the branches.
-            let free_rc_vars = {
+            // So we can make sure they are consumed in the current environment and all branches.
-                let mut free_rc_vars = FreeRcVariables::default();
+            let consume_variables = {
+                let branch_envs = {
+                    let mut branch_environments =
+                        Vec::with_capacity_in(new_branches.len() + 1, arena);
 
-                new_branches.iter().for_each(|(_, _, _, free_variables)| {
+                    for (_, _, _, branch_env) in new_branches.iter() {
-                    free_rc_vars.extend(free_variables);
+                        branch_environments.push(branch_env);
-                });
+                    }
 
-                free_rc_vars.extend(&new_default_branch.2);
+                    branch_environments.push(&new_default_branch.2);
 
-                // The scrutinee is not reference counted as it is a number.
+                    branch_environments
-                // Thus, we don't need to insert it in free vars.
+                };
-                // free_rc_vars.insert(*cond_symbol);
 
-                free_rc_vars
+                {
+                    let mut consume_variables = MutSet::default();
+
+                    for (symbol, ownership) in environment.variables_ownership.iter() {
+                        match ownership {
+                            Ownership::Owned => {
+                                if branch_envs
+                                .iter()
+                                .any(|branch_env| matches!(branch_env.get_variable_ownership(symbol).expect("All symbols defined in the current environment should be in the environment of the branches."), Ownership::Borrowed))
+                            {
+                                // If the variable is currently owned, and not in a some branches, it must be consumed in all branches
+                                consume_variables.insert(*symbol);
+                            }
+                                // Otherwise it can stay owned.
+                            }
+                            Ownership::Borrowed => {
+                                // If the variable is currently borrowed, it must be borrowed in all branches and we don't have to do anything.
+                            }
+                        }
+                    }
+
+                    consume_variables
+                }
             };
 
-            // If we compare the free variables from a branch with the free variables from all branches.
+            // Given the consume_variables we can determine what additional variables should be dropped in each branch.
-            // 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()
-                    .map(|(label, info, branch, free_variables)| {
+                    .map(|(label, info, branch, branch_env)| {
-                        let drop_vars = free_rc_vars.difference(&free_variables).collect();
+                        // If the variable is owned in the branch, it is not used in the branch and we can drop it.
-                        let newer_branch = insert_dec_stmts(arena, environment, &drop_vars, branch);
+                        let consume_variables_branch = consume_variables.iter().copied().filter(|consume_variable| {
+                             matches!(branch_env.get_variable_ownership(consume_variable).expect("All symbols defined in the current environment should be in the environment of the branches."), Ownership::Owned) 
+                        });
+
+                        let newer_branch = insert_dec_stmts(arena,  consume_variables_branch, branch);
                         (label, info, newer_branch.clone())
                     }),
                 arena,
@@ -549,33 +689,36 @@ fn insert_refcount_operations_stmt<'a>(
             .into_bump_slice();
 
             let newer_default_branch = {
-                let (info, branch, free_variables) = new_default_branch;
+                let (info, branch, branch_env) = new_default_branch;
-                let drop_vars = free_rc_vars.difference(&free_variables).collect();
+                // If the variable is owned in the branch, it is not used in the branch and we can drop it.
-                let newer_branch = insert_dec_stmts(arena, environment, &drop_vars, branch);
+                let consume_variables_branch = consume_variables.iter().copied().filter(|consume_variable| {
+                    matches!(branch_env.get_variable_ownership(consume_variable).expect("All symbols defined in the current environment should be in the environment of the branches."), Ownership::Owned) 
+                 });
+
+                let newer_branch = insert_dec_stmts(arena, consume_variables_branch, branch);
 
                 (info, newer_branch)
             };
 
-            let new_switch = arena.alloc(Stmt::Switch {
+            // In addition to updating the branches, we need to update the current environment.
+            for consume_variable in consume_variables.iter() {
+                environment.consume_variable(consume_variable);
+            }
+
+            arena.alloc(Stmt::Switch {
                 cond_symbol: *cond_symbol,
                 cond_layout: *cond_layout,
                 branches: newer_branches,
                 default_branch: newer_default_branch,
                 ret_layout: *ret_layout,
-            });
+            })
-
-            (new_switch, free_rc_vars)
         }
         Stmt::Ret(s) => {
-            // TODO use with_capacity if possible, or use an initializer to oneline this.
+            let ownership = environment.consume_variable(s);
-            let mut free_variables = FreeRcVariables::with_capacity_and_hasher(
+            debug_assert!(matches!(ownership, None | Some(Ownership::Owned))); // the return value should be owned or not reference counted at the return.
-                1,
+            return arena.alloc(Stmt::Ret(*s));
-                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::Refcounting(_, _) => unreachable!("refcounting should not be in the AST yet"),
         Stmt::Expect {
             condition,
             region,
@@ -600,39 +743,124 @@ fn insert_refcount_operations_stmt<'a>(
             parameters,
             body,
             remainder,
-        } => todo!(),
+        } => {
-        Stmt::Jump(_, _) => todo!(),
+            // Assuming that the values in the closure of the body of this jointpoint are already bound.
+
+            // Assuming that all variables are still owned. (So that we can determine what variables got consumed in the join point.)
+            debug_assert!(environment
+                .variables_ownership
+                .iter()
+                .all(|(_, ownership)| matches!(ownership, Ownership::Owned)));
+
+            let mut join_point_env = environment.clone();
+
+            let mut parameter_variables = parameters.iter().map(|Param { symbol, .. }| *symbol);
+
+            join_point_env.add_variables(parameter_variables.by_ref());
+
+            let new_body = insert_refcount_operations_stmt(arena, &mut join_point_env, body);
+
+            // We save the parameters consumed by this join point. So we can do the same when we jump to this joinpoint.
+            // This includes parameter variables, this might help with unused closure variables.
+            let consumed_variables = {
+                let consumed_variables = join_point_env
+                    .variables_ownership
+                    .iter()
+                    .filter_map(|(symbol, ownership)| match ownership {
+                        Ownership::Borrowed => Some(*symbol),
+                        _ => None,
+                    })
+                    .collect::<MutSet<_>>();
+
+                let consumed_closure = consumed_variables
+                    .difference(&parameter_variables.by_ref().collect::<MutSet<_>>())
+                    .copied()
+                    .collect::<MutSet<Symbol>>();
+
+                let consumed_parameters = Vec::from_iter_in(
+                    parameter_variables.map(|parameter| {
+                        if consumed_closure.contains(&parameter) {
+                            Consumption::Consumed
+                        } else {
+                            Consumption::Unconsumed
+                        }
+                    }),
+                    arena,
+                )
+                .into_bump_slice();
+
+                JoinPointConsumption {
+                    closure: consumed_closure,
+                    parameters: consumed_parameters,
+                }
+            };
+
+            environment.add_joinpoint_consumption(*id, consumed_variables);
+            let new_remainder = insert_refcount_operations_stmt(arena, environment, remainder);
+            environment.remove_joinpoint_consumption(*id);
+
+            arena.alloc(Stmt::Join {
+                id: *id,
+                parameters: parameters.clone(),
+                body: new_body,
+                remainder: new_remainder,
+            })
+        }
+        Stmt::Jump(join_point_id, arguments) => {
+            let JoinPointConsumption {
+                closure: consumed_variables,
+                parameters: consumed_parameters,
+            } = environment.get_joinpoint_consume(*join_point_id);
+
+            // the consumed variables contain the arguments, so we don't need ot
+            for consumed_variable in consumed_variables.iter() {
+                environment.consume_variable(consumed_variable);
+            }
+
+            // consume all arguments that are consumed by the join point.
+            for (argument, _) in arguments
+                .iter()
+                .zip(consumed_parameters.iter())
+                .filter(|(_, consumption)| matches!(consumption, Consumption::Consumed))
+            {
+                let ownership = environment.consume_variable(argument);
+                // the argument should be owned or not reference counted at the return.
+                debug_assert!(matches!(ownership, None | Some(Ownership::Owned)));
+            }
+
+            arena.alloc(Stmt::Jump(*join_point_id, arguments.clone()))
+        }
         Stmt::Crash(_, _) => todo!(),
     }
 }
 
 /**
- * Insert increment statements for the given symbols compensating for the ownership.
+Insert increment statements for the given symbols compensating for the ownership.
- */
+*/
-fn insert_inc_stmts<'a>(
+fn consume_and_insert_inc_stmts<'a>(
     arena: &'a Bump,
     environment: &mut Environment,
-    usage: &VariableUsage,
+    usage: &OwnedUsage,
     continuation: &'a Stmt<'a>,
 ) -> &'a Stmt<'a> {
     usage
         .iter()
         .fold(continuation, |continuation, (symbol, usage_count)| {
-            insert_inc_stmt(arena, environment, symbol, *usage_count, continuation)
+            consume_and_insert_inc_stmt(arena, environment, symbol, *usage_count, continuation)
         })
 }
 
 /**
 Insert an increment statement for the given symbol compensating for the ownership.
 */
-fn insert_inc_stmt<'a>(
+fn consume_and_insert_inc_stmt<'a>(
     arena: &'a Bump,
     environment: &mut Environment,
     symbol: &Symbol,
     usage_count: u64,
     continuation: &'a Stmt<'a>,
 ) -> &'a Stmt<'a> {
-    let new_count = match environment.consume_variable(symbol) {
+    let new_count = match environment.consume_rc_variable(symbol) {
         // If the variable is borrowed, we need to increment the reference count for each usage.
         Ownership::Borrowed => usage_count,
         // If the variable is owned, we need to increment the reference count for each usage except one.
@@ -651,30 +879,53 @@ fn insert_inc_stmt<'a>(
 /**
 Insert decrement statements for the given symbols if they are owned.
 */
-fn insert_dec_stmts<'a>(
+fn consume_and_insert_dec_stmts<'a>(
     arena: &'a Bump,
     environment: &mut Environment,
-    symbols: &MutSet<&Symbol>,
+    symbols: impl Iterator<Item = &'a Symbol>,
     continuation: &'a Stmt<'a>,
 ) -> &'a Stmt<'a> {
-    symbols.iter().fold(continuation, |continuation, symbol| {
+    symbols.fold(continuation, |continuation, symbol| {
-        insert_dec_stmt(arena, environment, symbol, continuation)
+        consume_and_insert_dec_stmt(arena, environment, symbol, continuation)
     })
 }
 
 /**
 Insert a decrement statement for the given symbol if it is owned.
 */
-fn insert_dec_stmt<'a>(
+fn consume_and_insert_dec_stmt<'a>(
     arena: &'a Bump,
     environment: &mut Environment,
     symbol: &Symbol,
     continuation: &'a Stmt<'a>,
 ) -> &'a Stmt<'a> {
-    match environment.consume_variable(symbol) {
+    match environment.consume_rc_variable(symbol) {
         // If the variable is borrowed, don't have to decrement the reference count.
         Ownership::Borrowed => continuation,
         // If the variable is owned, we do need to decrement the reference count.
-        Ownership::Owned => arena.alloc(Stmt::Refcounting(ModifyRc::Dec(*symbol), continuation)),
+        Ownership::Owned => insert_dec_stmt(arena, *symbol, continuation),
     }
 }
+
+/**
+Insert decrement statements for the given symbols.
+*/
+fn insert_dec_stmts<'a, 's>(
+    arena: &'a Bump,
+    symbols: impl Iterator<Item = Symbol>,
+    continuation: &'a Stmt<'a>,
+) -> &'a Stmt<'a> {
+    symbols.fold(continuation, |continuation, symbol| {
+        insert_dec_stmt(arena, symbol, continuation)
+    })
+}
+/**
+Insert a decrement statement for the given symbol.
+*/
+fn insert_dec_stmt<'a, 's>(
+    arena: &'a Bump,
+    symbol: Symbol,
+    continuation: &'a Stmt<'a>,
+) -> &'a Stmt<'a> {
+    arena.alloc(Stmt::Refcounting(ModifyRc::Dec(symbol), continuation))
+}