mirror of
https://github.com/roc-lang/roc.git
synced 2025-08-02 19:32:17 +00:00
Start new free var inference and joint points
This commit is contained in:
parent
9c73b5041a
commit
0afe94e6db
1 changed files with 468 additions and 217 deletions
|
@ -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();
|
||||
variable_rc_types_env.insert_variables_rc_type_proc(&proc);
|
||||
let variable_rc_types_env = variable_rc_types_env.clone();
|
||||
|
||||
let mut initial_environment = Environment {
|
||||
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();
|
||||
insert_refcount_operations_proc(arena, variable_rc_types_env, proc);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
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>;
|
||||
|
||||
// Combine variable usage by summing the usage of each variable.
|
||||
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 map keeping track of how many times a variable is used as owned.
|
||||
*/
|
||||
type OwnedUsage = MutMap<Symbol, u64>;
|
||||
|
||||
/**
|
||||
A struct to determine the usage of variables in an expression.
|
||||
*/
|
||||
struct VariableUsageEnv<'a> {
|
||||
variable_usage: VariableUsage,
|
||||
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.
|
||||
*/
|
||||
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::Tag { arguments, .. } | Expr::Struct(arguments) => {
|
||||
usage_env.insert_variable_usages(arguments.iter().copied());
|
||||
|
||||
Expr::Call(Call { arguments, .. })
|
||||
| Expr::Tag { arguments, .. }
|
||||
| Expr::Struct(arguments) => {
|
||||
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(
|
||||
|element| match element {
|
||||
Self::owned_usages(
|
||||
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.
|
||||
If the symbol is not reference counted, it will be ignored.
|
||||
Return owned usages.
|
||||
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) {
|
||||
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,
|
||||
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>,
|
||||
fn owned_usages(
|
||||
variable_rc_types: &VariableRcTypes,
|
||||
symbols: impl Iterator<Item = Symbol>,
|
||||
) {
|
||||
symbols.for_each(|symbol| self.insert_variable_usage(symbol));
|
||||
) -> VariableUsage {
|
||||
// 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 {
|
||||
// This function should only be called on reference counted variables.
|
||||
debug_assert!(matches!(
|
||||
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"),
|
||||
fn consume_variable(self: &mut Environment<'v, 'a>, variable: &Symbol) -> Option<Ownership> {
|
||||
if !self.variables_ownership.contains_key(variable) {
|
||||
return None;
|
||||
}
|
||||
|
||||
// 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) {
|
||||
match self.get_variable_rc_type(variable) {
|
||||
fn add_variable(self: &mut Environment<'v, 'a>, variable: Symbol) {
|
||||
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 the variable afterwards to prevent it from being used outside the function call.
|
||||
/**
|
||||
Remove variables from the environment.
|
||||
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
|
||||
where
|
||||
F: Fn(&mut Environment) -> R,
|
||||
{
|
||||
match self.get_variable_rc_type(variable) {
|
||||
VarRcType::ReferenceCounted => {
|
||||
self.variables_ownership.insert(*variable, Ownership::Owned);
|
||||
let result = callback(self);
|
||||
self.variables_ownership.remove(&variable);
|
||||
result
|
||||
}
|
||||
VarRcType::NotReferenceCounted => callback(self),
|
||||
}
|
||||
fn add_joinpoint_consumption(
|
||||
self: &mut Environment<'v, 'a>,
|
||||
joinpoint_id: JoinPointId,
|
||||
consumption: JoinPointConsumption<'a>,
|
||||
) {
|
||||
self.jointpoint_closures.insert(joinpoint_id, consumption);
|
||||
}
|
||||
|
||||
/**
|
||||
Get the consumed closure from a join point id.
|
||||
*/
|
||||
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.
|
||||
let (new_stmt, mut free_rc_vars_stmt) = environment.with_variable(binding, |env| {
|
||||
insert_refcount_operations_stmt(arena, env, stmt)
|
||||
});
|
||||
// First evaluate the continuation and let it consume it's free variables.
|
||||
environment.add_variable(*binding); // Add the bound variable to the environment. As it can be used in the continuation.
|
||||
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.
|
||||
let new_let_with_refcount =
|
||||
insert_inc_stmts(arena, environment, &variable_usage, new_let);
|
||||
// Insert the reference count operations for the owned variables used in the expression.
|
||||
consume_and_insert_inc_stmts(arena, environment, &owned_usage, new_let)
|
||||
}
|
||||
VariableUsage::Borrowed(symbol) => {
|
||||
let newer_stmt =
|
||||
consume_and_insert_dec_stmt(arena, environment, &symbol, new_stmt);
|
||||
|
||||
let free_rc_vars = {
|
||||
// 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)
|
||||
arena.alloc(Stmt::Let(*binding, expr.clone(), *layout, newer_stmt))
|
||||
}
|
||||
}
|
||||
}
|
||||
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) =
|
||||
insert_refcount_operations_stmt(arena, &mut environment.clone(), branch);
|
||||
let mut branch_env = environment.clone();
|
||||
|
||||
(*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.
|
||||
let free_rc_vars = {
|
||||
let mut free_rc_vars = FreeRcVariables::default();
|
||||
// Determine what variables are consumed in some of the branches.
|
||||
// So we can make sure they are consumed in the current environment and all branches.
|
||||
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)| {
|
||||
free_rc_vars.extend(free_variables);
|
||||
});
|
||||
for (_, _, _, branch_env) in new_branches.iter() {
|
||||
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.
|
||||
// Thus, we don't need to insert it in free vars.
|
||||
// free_rc_vars.insert(*cond_symbol);
|
||||
branch_environments
|
||||
};
|
||||
|
||||
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.
|
||||
// We can determine which variables are not used in that branch.
|
||||
// And as such, we can drop them before continuing.
|
||||
// Given the consume_variables we can determine what additional variables should be dropped in each branch.
|
||||
let newer_branches = Vec::from_iter_in(
|
||||
new_branches
|
||||
.into_iter()
|
||||
.map(|(label, info, branch, free_variables)| {
|
||||
let drop_vars = free_rc_vars.difference(&free_variables).collect();
|
||||
let newer_branch = insert_dec_stmts(arena, environment, &drop_vars, branch);
|
||||
.map(|(label, info, branch, branch_env)| {
|
||||
// If the variable is owned in the branch, it is not used in the branch and we can drop it.
|
||||
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 drop_vars = free_rc_vars.difference(&free_variables).collect();
|
||||
let newer_branch = insert_dec_stmts(arena, environment, &drop_vars, branch);
|
||||
let (info, branch, branch_env) = new_default_branch;
|
||||
// If the variable is owned in the branch, it is not used in the branch and we can drop it.
|
||||
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 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);
|
||||
let ownership = environment.consume_variable(s);
|
||||
debug_assert!(matches!(ownership, None | Some(Ownership::Owned))); // the return value should be owned or not reference counted at the return.
|
||||
return arena.alloc(Stmt::Ret(*s));
|
||||
}
|
||||
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(¶meter_variables.by_ref().collect::<MutSet<_>>())
|
||||
.copied()
|
||||
.collect::<MutSet<Symbol>>();
|
||||
|
||||
let consumed_parameters = Vec::from_iter_in(
|
||||
parameter_variables.map(|parameter| {
|
||||
if consumed_closure.contains(¶meter) {
|
||||
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.
|
||||
*/
|
||||
fn insert_inc_stmts<'a>(
|
||||
Insert increment statements for the given symbols compensating for the ownership.
|
||||
*/
|
||||
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| {
|
||||
insert_dec_stmt(arena, environment, symbol, continuation)
|
||||
symbols.fold(continuation, |continuation, symbol| {
|
||||
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))
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue