From ae9de11e77d512a477fc93e4585580a1d725136c Mon Sep 17 00:00:00 2001 From: microproofs Date: Thu, 20 Jul 2023 21:02:33 -0400 Subject: [PATCH] big checkpoint: feat: add monomorphize and other useful tree function abstractions feat: started testing function hositing result so far --- crates/aiken-lang/src/gen_uplc.rs | 1 - crates/aiken-lang/src/gen_uplc2.rs | 776 +++++------------- crates/aiken-lang/src/gen_uplc2/air.rs | 4 +- crates/aiken-lang/src/gen_uplc2/builder.rs | 526 +++--------- crates/aiken-lang/src/gen_uplc2/tree.rs | 890 ++++++++++++++++++++- 5 files changed, 1193 insertions(+), 1004 deletions(-) diff --git a/crates/aiken-lang/src/gen_uplc.rs b/crates/aiken-lang/src/gen_uplc.rs index 0fbf4499..948173a1 100644 --- a/crates/aiken-lang/src/gen_uplc.rs +++ b/crates/aiken-lang/src/gen_uplc.rs @@ -103,7 +103,6 @@ impl<'a> CodeGenerator<'a> { FunctionAccessKey { module_name, function_name, - variant_name, }, value, ) diff --git a/crates/aiken-lang/src/gen_uplc2.rs b/crates/aiken-lang/src/gen_uplc2.rs index 0c0f0c39..2ecf7b97 100644 --- a/crates/aiken-lang/src/gen_uplc2.rs +++ b/crates/aiken-lang/src/gen_uplc2.rs @@ -18,16 +18,13 @@ use crate::{ }, builtins::{bool, data, int, void}, expr::TypedExpr, - gen_uplc::{ - air::Air, - builder::{ - self as build, get_arg_type_name, AssignmentProperties, ClauseProperties, DataTypeKey, - FunctionAccessKey, SpecificClause, - }, + gen_uplc::builder::{ + self as build, get_arg_type_name, AssignmentProperties, ClauseProperties, DataTypeKey, + FunctionAccessKey, SpecificClause, }, gen_uplc2::builder::{ - convert_opaque_type, erase_opaque_operations, find_and_replace_generics, - get_generic_id_and_type, + convert_opaque_type, erase_opaque_type_operations, find_and_replace_generics, + get_generic_id_and_type, get_variant_name, monomorphize, }, tipo::{ ModuleValueConstructor, PatternConstructor, Type, TypeInfo, ValueConstructor, @@ -36,22 +33,11 @@ use crate::{ }; use self::{ - builder::{IndexCounter, TreePath}, - tree::{AirExpression, AirStatement, AirTree}, + air::Air, + builder::{CodeGenFunction, UserFunction}, + tree::{AirExpression, AirTree, TreePath}, }; -#[derive(Clone, Debug)] -pub enum CodeGenFunction { - Function(AirTree, Vec), - Link(String), -} - -#[derive(Clone, Debug)] -pub enum UserFunction { - Function(AirTree, Term), - Link(String), -} - #[derive(Clone)] pub struct CodeGenerator<'a> { defined_functions: IndexMap, @@ -103,10 +89,9 @@ impl<'a> CodeGenerator<'a> { let mut validator_args_tree = self.check_validator_args(&fun.arguments, true, air_tree_fun); validator_args_tree = AirTree::no_op().hoist_over(validator_args_tree); - println!("{:#?}", validator_args_tree); println!("{:#?}", validator_args_tree.to_vec()); - let full_tree = self.hoist_functions(validator_args_tree); + let full_tree = self.hoist_functions_to_validator(validator_args_tree); todo!() } @@ -115,10 +100,9 @@ impl<'a> CodeGenerator<'a> { let mut air_tree = self.build(test_body); air_tree = AirTree::no_op().hoist_over(air_tree); - println!("{:#?}", air_tree); println!("{:#?}", air_tree.to_vec()); - let full_tree = self.hoist_functions(air_tree); + let full_tree = self.hoist_functions_to_validator(air_tree); todo!() } @@ -291,7 +275,13 @@ impl<'a> CodeGenerator<'a> { right, tipo, .. - } => AirTree::binop(*name, tipo.clone(), self.build(left), self.build(right)), + } => AirTree::binop( + *name, + tipo.clone(), + self.build(left), + self.build(right), + left.tipo(), + ), TypedExpr::Assignment { tipo, @@ -388,6 +378,7 @@ impl<'a> CodeGenerator<'a> { let when_assign = AirTree::when( subject_name, tipo.clone(), + subject.tipo(), AirTree::local_var(constr_var, subject.tipo()), clauses, ); @@ -553,9 +544,10 @@ impl<'a> CodeGenerator<'a> { let expect = AirTree::binop( BinOp::Eq, - int(), + bool(), AirTree::int(expected_int), AirTree::local_var(name, int()), + int(), ); AirTree::assert_bool(true, assignment.hoist_over(expect)) } else { @@ -1279,6 +1271,7 @@ impl<'a> CodeGenerator<'a> { let when_expr = AirTree::when( format!("__subject_span_{}_{}", location.start, location.end), + void(), tipo.clone(), AirTree::local_var( format!("__constr_var_span_{}_{}", location.start, location.end), @@ -1953,7 +1946,7 @@ impl<'a> CodeGenerator<'a> { name_index_assigns.iter().for_each(|(name, index, _)| { if let Some((index, prev_name)) = defined_indices .iter() - .find(|(defined_index, nm)| defined_index == index) + .find(|(defined_index, _nm)| defined_index == index) { previous_defined_names.push((*index, prev_name.clone(), name.clone())); } else if name != "_" { @@ -2210,575 +2203,216 @@ impl<'a> CodeGenerator<'a> { AirTree::anon_func(arg_names, arg_assigns.hoist_over(body)) } - fn hoist_functions(&mut self, mut air_tree: AirTree) -> AirTree { + fn hoist_functions_to_validator(&mut self, mut air_tree: AirTree) -> AirTree { let mut functions_to_hoist = IndexMap::new(); - let mut function_holder = IndexMap::new(); + let mut used_functions = vec![]; + let mut defined_functions = vec![]; - erase_opaque_operations(&mut air_tree, &self.data_types); + erase_opaque_type_operations(&mut air_tree, &self.data_types); self.find_function_vars_and_depth( - &air_tree, + &mut air_tree, &mut functions_to_hoist, - &mut function_holder, - &mut TreePath::new(), - 0, - 0, + &mut used_functions, ); - todo!() + println!("FUNCTIONS TO HOIST {:#?}", functions_to_hoist); + + while let Some((key, variant_name)) = used_functions.pop() { + defined_functions.push((key.clone(), variant_name.clone())); + let function_variants = functions_to_hoist + .get(&key) + .unwrap_or_else(|| panic!("Missing Function Definition")); + + let (_, function) = function_variants + .get(&variant_name) + .unwrap_or_else(|| panic!("Missing Function Variant Definition")); + + if let UserFunction::Function(body, deps) = function { + let mut hoist_body = body.clone(); + let mut hoist_deps = deps.clone(); + + self.hoist_functions( + &mut hoist_body, + &mut functions_to_hoist, + &mut used_functions, + &defined_functions, + &mut hoist_deps, + ); + + let function_variants = functions_to_hoist + .get_mut(&key) + .unwrap_or_else(|| panic!("Missing Function Definition")); + + let (_, function) = function_variants + .get_mut(&variant_name) + .unwrap_or_else(|| panic!("Missing Function Variant Definition")); + + *function = UserFunction::Function(hoist_body, hoist_deps); + } else { + todo!("Deal with Link later") + } + } + + println!("FUNCTIONS TO HOIST {:#?}", functions_to_hoist); + println!("FUNCTIONS DEFINED {:#?}", defined_functions); + println!("FUNCTIONS USED {:#?}", used_functions); + + air_tree + } + + fn hoist_functions( + &mut self, + air_tree: &mut AirTree, + function_usage: &mut IndexMap< + FunctionAccessKey, + IndexMap, + >, + used_functions: &mut Vec<(FunctionAccessKey, String)>, + defined_functions: &[(FunctionAccessKey, String)], + current_function_deps: &mut Vec<(FunctionAccessKey, String)>, + ) { + self.find_function_vars_and_depth(air_tree, function_usage, current_function_deps); + + for (generic_function_key, variant_name) in current_function_deps.iter() { + if !used_functions + .iter() + .any(|(key, name)| key == generic_function_key && name == variant_name) + && !defined_functions + .iter() + .any(|(key, name)| key == generic_function_key && name == variant_name) + { + used_functions.push((generic_function_key.clone(), variant_name.clone())); + } + } } fn find_function_vars_and_depth( &mut self, - validator_args_tree: &AirTree, - function_usage: &mut IndexMap>, - function_holder: &mut IndexMap>, - tree_path: &mut TreePath, - current_depth: usize, - depth_index: usize, + air_tree: &mut AirTree, + function_usage: &mut IndexMap< + FunctionAccessKey, + IndexMap, + >, + dependency_functions: &mut Vec<(FunctionAccessKey, String)>, ) { - let mut index_count = IndexCounter::new(); - tree_path.push(current_depth, depth_index); - match validator_args_tree { - AirTree::Statement { - statement, - hoisted_over: Some(hoisted_over), - } => { - match statement { - AirStatement::Let { value, .. } => { - self.find_function_vars_and_depth( - value, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirStatement::DefineFunc { func_body, .. } => { - self.find_function_vars_and_depth( - func_body, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirStatement::AssertConstr { constr, .. } => { - self.find_function_vars_and_depth( - constr, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirStatement::AssertBool { value, .. } => { - self.find_function_vars_and_depth( - value, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirStatement::ClauseGuard { pattern, .. } => { - self.find_function_vars_and_depth( - pattern, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirStatement::ListClauseGuard { .. } => {} - AirStatement::TupleGuard { .. } => {} - AirStatement::FieldsExpose { record, .. } => { - self.find_function_vars_and_depth( - record, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirStatement::ListAccessor { list, .. } => { - self.find_function_vars_and_depth( - list, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirStatement::ListExpose { .. } => {} - AirStatement::TupleAccessor { tuple, .. } => { - self.find_function_vars_and_depth( - tuple, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirStatement::NoOp => {} + air_tree.traverse_tree_with(&mut |air_tree, tree_path| { + if let AirTree::Expression(AirExpression::Var { constructor, .. }) = air_tree { + let ValueConstructorVariant::ModuleFn { + name: func_name, + module, + builtin: None, + .. + } = &constructor.variant + else { return }; + + let function_var_tipo = &constructor.tipo; + println!("FUNCTION VAR TYPE {:#?}", function_var_tipo); + + let generic_function_key = FunctionAccessKey { + module_name: module.clone(), + function_name: func_name.clone(), }; - self.find_function_vars_and_depth( - hoisted_over, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), + let function_def = self + .functions + .get(&generic_function_key) + .unwrap_or_else(|| panic!("Missing Function Definition")); + + println!("Function Def {:#?}", function_def); + + let mut function_var_types = function_var_tipo + .arg_types() + .unwrap_or_else(|| panic!("Expected a function tipo with arg types")); + + function_var_types.push( + function_var_tipo + .return_type() + .unwrap_or_else(|| panic!("Should have return type")), ); - } - AirTree::Expression(e) => match e { - AirExpression::List { items, .. } => { - for item in items { - self.find_function_vars_and_depth( - item, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } + + let mut function_def_types = function_def + .arguments + .iter() + .map(|arg| &arg.tipo) + .collect_vec(); + + function_def_types.push(&function_def.return_type); + + let mono_types: IndexMap> = if !function_def_types.is_empty() { + function_def_types + .into_iter() + .zip(function_var_types.into_iter()) + .flat_map(|(func_tipo, var_tipo)| { + get_generic_id_and_type(func_tipo, &var_tipo) + }) + .collect() + } else { + IndexMap::new() + }; + + println!("MONO TYPES {:#?}", mono_types); + + let variant_name = mono_types + .iter() + .sorted_by(|(id, _), (id2, _)| id.cmp(id2)) + .map(|(_, tipo)| get_variant_name(tipo)) + .join(""); + + if !dependency_functions + .iter() + .any(|(key, name)| key == &generic_function_key && name == &variant_name) + { + dependency_functions.push((generic_function_key.clone(), variant_name.clone())); } - AirExpression::Tuple { items, .. } => { - for item in items { - self.find_function_vars_and_depth( - item, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - } - AirExpression::Var { - constructor, - name, - variant_name, - } => { - let ValueConstructorVariant::ModuleFn { - name: func_name, - module, - builtin: None, - .. - } = &constructor.variant - else { return }; - let function_var_tipo = &constructor.tipo; - println!("FUNCTION VAR TYPE {:#?}", function_var_tipo); - - let generic_function_key = FunctionAccessKey { - module_name: module.clone(), - function_name: func_name.clone(), - }; - - let function_def = self - .functions - .get(&generic_function_key) - .unwrap_or_else(|| panic!("Missing Function Definition")); - - println!("Function Def {:#?}", function_def); - - let mut function_var_types = function_var_tipo - .arg_types() - .unwrap_or_else(|| panic!("Expected a function tipo with arg types")); - - function_var_types.push( - function_var_tipo - .return_type() - .unwrap_or_else(|| panic!("Should have return type")), - ); - - let mut function_def_types = function_def - .arguments - .iter() - .map(|arg| &arg.tipo) - .collect_vec(); - - function_def_types.push(&function_def.return_type); - - let mono_types: IndexMap> = if !function_def_types.is_empty() { - function_def_types - .into_iter() - .zip(function_var_types.into_iter()) - .flat_map(|(func_tipo, var_tipo)| { - get_generic_id_and_type(func_tipo, &var_tipo) - }) - .collect() + if let Some(func_variants) = function_usage.get_mut(&generic_function_key) { + if let Some((path, _)) = func_variants.get_mut(&variant_name) { + *path = path.common_ancestor(tree_path); } else { - IndexMap::new() - }; + let mut function_air_tree_body = self.build(&function_def.body); - todo!() - } - AirExpression::Call { func, args, .. } => { - self.find_function_vars_and_depth( - func, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); + monomorphize(&mut function_air_tree_body, &mono_types); - for arg in args { - self.find_function_vars_and_depth( - arg, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), + erase_opaque_type_operations(&mut function_air_tree_body, &self.data_types); + + func_variants.insert( + variant_name, + ( + tree_path.current_path(), + UserFunction::Function(function_air_tree_body, vec![]), + ), ); } - } - AirExpression::Fn { func_body, .. } => { - self.find_function_vars_and_depth( - func_body, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::Builtin { args, .. } => { - for arg in args { - self.find_function_vars_and_depth( - arg, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - } - AirExpression::BinOp { left, right, .. } => { - self.find_function_vars_and_depth( - left, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), + } else { + let mut function_air_tree_body = self.build(&function_def.body); + + monomorphize(&mut function_air_tree_body, &mono_types); + + erase_opaque_type_operations(&mut function_air_tree_body, &self.data_types); + + let mut function_variant_path = IndexMap::new(); + + function_variant_path.insert( + variant_name, + ( + tree_path.current_path(), + UserFunction::Function(function_air_tree_body, vec![]), + ), ); - self.find_function_vars_and_depth( - right, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); + function_usage.insert(generic_function_key, function_variant_path); } - AirExpression::UnOp { arg, .. } => { - self.find_function_vars_and_depth( - arg, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::UnWrapData { value, .. } => { - self.find_function_vars_and_depth( - value, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::WrapData { value, .. } => { - self.find_function_vars_and_depth( - value, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::When { - subject, clauses, .. - } => { - self.find_function_vars_and_depth( - subject, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); + } + }); + } - self.find_function_vars_and_depth( - clauses, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::Clause { - pattern, - then, - otherwise, - .. - } => { - self.find_function_vars_and_depth( - pattern, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); + fn uplc_code_gen(&mut self, ir_stack: &mut Vec) -> Term { + let mut arg_stack: Vec> = vec![]; - self.find_function_vars_and_depth( - then, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - - self.find_function_vars_and_depth( - otherwise, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::ListClause { - then, otherwise, .. - } => { - self.find_function_vars_and_depth( - then, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - - self.find_function_vars_and_depth( - otherwise, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::WrapClause { then, otherwise } => { - self.find_function_vars_and_depth( - then, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - - self.find_function_vars_and_depth( - otherwise, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::TupleClause { - then, otherwise, .. - } => { - self.find_function_vars_and_depth( - then, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - - self.find_function_vars_and_depth( - otherwise, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::Finally { pattern, then } => { - self.find_function_vars_and_depth( - pattern, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - - self.find_function_vars_and_depth( - then, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::If { - pattern, - then, - otherwise, - .. - } => { - self.find_function_vars_and_depth( - pattern, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - - self.find_function_vars_and_depth( - then, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - - self.find_function_vars_and_depth( - otherwise, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::Constr { args, .. } => { - for arg in args { - self.find_function_vars_and_depth( - arg, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - } - AirExpression::RecordUpdate { record, args, .. } => { - self.find_function_vars_and_depth( - record, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - for arg in args { - self.find_function_vars_and_depth( - arg, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - } - AirExpression::RecordAccess { record, .. } => { - self.find_function_vars_and_depth( - record, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::TupleIndex { tuple, .. } => { - self.find_function_vars_and_depth( - tuple, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::Trace { msg, then, .. } => { - self.find_function_vars_and_depth( - msg, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - - self.find_function_vars_and_depth( - then, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::FieldsEmpty { constr } => { - self.find_function_vars_and_depth( - constr, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - AirExpression::ListEmpty { list } => { - self.find_function_vars_and_depth( - list, - function_usage, - function_holder, - tree_path, - current_depth + 1, - index_count.next(), - ); - } - _ => {} - }, - _ => unreachable!(), + while let Some(ir_element) = ir_stack.pop() { + todo!() } - tree_path.pop(); + arg_stack[0].clone() } } diff --git a/crates/aiken-lang/src/gen_uplc2/air.rs b/crates/aiken-lang/src/gen_uplc2/air.rs index 3cb22f6c..a832f70e 100644 --- a/crates/aiken-lang/src/gen_uplc2/air.rs +++ b/crates/aiken-lang/src/gen_uplc2/air.rs @@ -61,6 +61,7 @@ pub enum Air { BinOp { name: BinOp, tipo: Arc, + argument_tipo: Arc, }, UnOp { op: UnOp, @@ -85,6 +86,7 @@ pub enum Air { When { tipo: Arc, subject_name: String, + subject_tipo: Arc, }, Clause { subject_tipo: Arc, @@ -180,5 +182,3 @@ pub enum Air { FieldsEmpty, ListEmpty, } - -impl Air {} diff --git a/crates/aiken-lang/src/gen_uplc2/builder.rs b/crates/aiken-lang/src/gen_uplc2/builder.rs index 290ce86f..2ca283cd 100644 --- a/crates/aiken-lang/src/gen_uplc2/builder.rs +++ b/crates/aiken-lang/src/gen_uplc2/builder.rs @@ -1,53 +1,34 @@ +use std::sync::Arc; + use indexmap::IndexMap; use crate::{ - ast::TypedDataType, + ast::{Function, TypedDataType, TypedFunction}, builtins::bool, gen_uplc::builder::{lookup_data_type_by_tipo, DataTypeKey, FunctionAccessKey}, tipo::{TypeVar, ValueConstructorVariant}, }; -use std::sync::Arc; use crate::{ ast::{BinOp, ClauseGuard, Constant, UnOp}, tipo::Type, }; -use super::tree::{AirExpression, AirStatement, AirTree}; +use super::{ + air::Air, + tree::{AirExpression, AirStatement, AirTree, TreePath}, +}; #[derive(Clone, Debug)] -pub struct TreePath { - path: Vec<(usize, usize)>, +pub enum CodeGenFunction { + Function(AirTree, Vec), + Link(String), } -impl TreePath { - pub fn new() -> Self { - TreePath { path: vec![] } - } - - pub fn push(&mut self, depth: usize, index: usize) { - self.path.push((depth, index)); - } - - pub fn pop(&mut self) { - self.path.pop(); - } -} - -pub struct IndexCounter { - current_index: usize, -} - -impl IndexCounter { - pub fn new() -> Self { - IndexCounter { current_index: 0 } - } - - pub fn next(&mut self) -> usize { - let current_index = self.current_index; - self.current_index += 1; - current_index - } +#[derive(Clone, Debug)] +pub enum UserFunction { + Function(AirTree, Vec<(FunctionAccessKey, String)>), + Link(String), } pub fn get_generic_id_and_type(tipo: &Type, param: &Type) -> Vec<(u64, Arc)> { @@ -194,6 +175,8 @@ pub fn find_and_replace_generics( new_args.push(arg); } + println!("SO ARGS ARE {:#?}", new_args); + let ret = find_and_replace_generics(ret, mono_types); let t = Type::Fn { @@ -245,63 +228,123 @@ pub fn handle_clause_guard(clause_guard: &ClauseGuard>) -> AirTree { AirTree::unop(UnOp::Not, val) } ClauseGuard::Equals { left, right, .. } => { - let left = handle_clause_guard(left); - let right = handle_clause_guard(right); + let left_child = handle_clause_guard(left); + let right_child = handle_clause_guard(right); - AirTree::binop(BinOp::Eq, bool(), left, right) + AirTree::binop(BinOp::Eq, bool(), left_child, right_child, left.tipo()) } ClauseGuard::NotEquals { left, right, .. } => { - let left = handle_clause_guard(left); - let right = handle_clause_guard(right); + let left_child = handle_clause_guard(left); + let right_child = handle_clause_guard(right); - AirTree::binop(BinOp::NotEq, bool(), left, right) + AirTree::binop(BinOp::NotEq, bool(), left_child, right_child, left.tipo()) } ClauseGuard::GtInt { left, right, .. } => { - let left = handle_clause_guard(left); - let right = handle_clause_guard(right); + let left_child = handle_clause_guard(left); + let right_child = handle_clause_guard(right); - AirTree::binop(BinOp::GtInt, bool(), left, right) + AirTree::binop(BinOp::GtInt, bool(), left_child, right_child, left.tipo()) } ClauseGuard::GtEqInt { left, right, .. } => { - let left = handle_clause_guard(left); - let right = handle_clause_guard(right); + let left_child = handle_clause_guard(left); + let right_child = handle_clause_guard(right); - AirTree::binop(BinOp::GtEqInt, bool(), left, right) + AirTree::binop(BinOp::GtEqInt, bool(), left_child, right_child, left.tipo()) } ClauseGuard::LtInt { left, right, .. } => { - let left = handle_clause_guard(left); - let right = handle_clause_guard(right); + let left_child = handle_clause_guard(left); + let right_child = handle_clause_guard(right); - AirTree::binop(BinOp::LtInt, bool(), left, right) + AirTree::binop(BinOp::LtInt, bool(), left_child, right_child, left.tipo()) } ClauseGuard::LtEqInt { left, right, .. } => { - let left = handle_clause_guard(left); - let right = handle_clause_guard(right); + let left_child = handle_clause_guard(left); + let right_child = handle_clause_guard(right); - AirTree::binop(BinOp::LtEqInt, bool(), left, right) + AirTree::binop(BinOp::LtEqInt, bool(), left_child, right_child, left.tipo()) } ClauseGuard::Or { left, right, .. } => { - let left = handle_clause_guard(left); - let right = handle_clause_guard(right); + let left_child = handle_clause_guard(left); + let right_child = handle_clause_guard(right); - AirTree::binop(BinOp::Or, bool(), left, right) + AirTree::binop(BinOp::Or, bool(), left_child, right_child, left.tipo()) } ClauseGuard::And { left, right, .. } => { - let left = handle_clause_guard(left); - let right = handle_clause_guard(right); + let left_child = handle_clause_guard(left); + let right_child = handle_clause_guard(right); - AirTree::binop(BinOp::And, bool(), left, right) + AirTree::binop(BinOp::And, bool(), left_child, right_child, left.tipo()) } ClauseGuard::Var { tipo, name, .. } => AirTree::local_var(name, tipo.clone()), ClauseGuard::Constant(constant) => constants_ir(constant), } } -pub fn erase_opaque_operations( +pub fn get_variant_name(t: &Arc) -> String { + if t.is_string() { + "_string".to_string() + } else if t.is_int() { + "_int".to_string() + } else if t.is_bool() { + "_bool".to_string() + } else if t.is_bytearray() { + "_bytearray".to_string() + } else if t.is_map() { + let mut full_type = vec!["_map".to_string()]; + let pair_type = &t.get_inner_types()[0]; + let fst_type = &pair_type.get_inner_types()[0]; + let snd_type = &pair_type.get_inner_types()[1]; + full_type.push(get_variant_name(fst_type)); + full_type.push(get_variant_name(snd_type)); + full_type.join("") + } else if t.is_list() { + let full_type = "_list".to_string(); + let list_type = &t.get_inner_types()[0]; + + format!("{}{}", full_type, get_variant_name(list_type)) + } else if t.is_tuple() { + let mut full_type = vec!["_tuple".to_string()]; + + let inner_types = t.get_inner_types(); + + for arg_type in inner_types { + full_type.push(get_variant_name(&arg_type)); + } + full_type.join("") + } else if t.is_unbound() { + "_unbound".to_string() + } else { + let full_type = "_data".to_string(); + + if t.is_generic() { + println!("FULL TYPE: {:#?}", t); + panic!("FOUND A POLYMORPHIC TYPE. EXPECTED MONOMORPHIC TYPE"); + } + + full_type + } +} + +pub fn monomorphize(air_tree: &mut AirTree, mono_types: &IndexMap>) { + air_tree.traverse_tree_with(&mut |air_tree: &mut AirTree, _| { + let mut held_types = air_tree.mut_held_types(); + println!("Held types: {:#?}", held_types); + + while let Some(tipo) = held_types.pop() { + *tipo = find_and_replace_generics(tipo, mono_types) + } + }); +} + +pub fn function_deps(air_tree: &mut AirTree, mono_types: &IndexMap>) { + air_tree.traverse_tree_with(&mut |air_tree: &mut AirTree, _| {}); +} + +pub fn erase_opaque_type_operations( air_tree: &mut AirTree, data_types: &IndexMap, ) { - traverse_tree_with(air_tree, &mut TreePath::new(), 0, 0, &|air_tree, _| { + air_tree.traverse_tree_with(&mut |air_tree, _| { if let AirTree::Expression(e) = air_tree { match e { AirExpression::Constr { tipo, args, .. } => { @@ -325,372 +368,11 @@ pub fn erase_opaque_operations( hoisted_over: Some(hoisted_over), } = air_tree { - let name = indices[0].1.clone(); - if check_replaceable_opaque_type(&record.get_type(), data_types) { + if check_replaceable_opaque_type(&record.return_type(), data_types) { + let name = indices[0].1.clone(); *air_tree = AirTree::let_assignment(name, (**record).clone()) .hoist_over((**hoisted_over).clone()) } } }); } - -fn traverse_tree_with( - air_tree: &mut AirTree, - tree_path: &mut TreePath, - current_depth: usize, - depth_index: usize, - with: &impl Fn(&mut AirTree, &TreePath), -) { - let mut index_count = IndexCounter::new(); - tree_path.push(current_depth, depth_index); - match air_tree { - AirTree::Statement { - statement, - hoisted_over: Some(hoisted_over), - } => { - match statement { - AirStatement::Let { value, .. } => { - traverse_tree_with( - value, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirStatement::DefineFunc { func_body, .. } => { - traverse_tree_with( - func_body, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirStatement::AssertConstr { constr, .. } => { - traverse_tree_with( - constr, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirStatement::AssertBool { value, .. } => { - traverse_tree_with( - value, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirStatement::ClauseGuard { pattern, .. } => { - traverse_tree_with( - pattern, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirStatement::ListClauseGuard { .. } => {} - AirStatement::TupleGuard { .. } => {} - AirStatement::FieldsExpose { record, .. } => { - traverse_tree_with( - record, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirStatement::ListAccessor { list, .. } => { - traverse_tree_with( - list, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirStatement::ListExpose { .. } => {} - AirStatement::TupleAccessor { tuple, .. } => { - traverse_tree_with( - tuple, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirStatement::NoOp => {} - }; - - traverse_tree_with( - hoisted_over, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirTree::Expression(e) => match e { - AirExpression::List { items, .. } => { - for item in items { - traverse_tree_with( - item, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - } - AirExpression::Tuple { items, .. } => { - for item in items { - traverse_tree_with( - item, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - } - AirExpression::Var { - constructor, - name, - variant_name, - } => { - todo!() - } - AirExpression::Call { func, args, .. } => { - traverse_tree_with(func, tree_path, current_depth + 1, index_count.next(), with); - - for arg in args { - traverse_tree_with(arg, tree_path, current_depth + 1, index_count.next(), with); - } - } - AirExpression::Fn { func_body, .. } => { - traverse_tree_with( - func_body, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::Builtin { args, .. } => { - for arg in args { - traverse_tree_with(arg, tree_path, current_depth + 1, index_count.next(), with); - } - } - AirExpression::BinOp { left, right, .. } => { - traverse_tree_with(left, tree_path, current_depth + 1, index_count.next(), with); - - traverse_tree_with( - right, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::UnOp { arg, .. } => { - traverse_tree_with(arg, tree_path, current_depth + 1, index_count.next(), with); - } - AirExpression::UnWrapData { value, .. } => { - traverse_tree_with( - value, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::WrapData { value, .. } => { - traverse_tree_with( - value, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::When { - subject, clauses, .. - } => { - traverse_tree_with( - subject, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - - traverse_tree_with( - clauses, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::Clause { - pattern, - then, - otherwise, - .. - } => { - traverse_tree_with( - pattern, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - - traverse_tree_with(then, tree_path, current_depth + 1, index_count.next(), with); - - traverse_tree_with( - otherwise, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::ListClause { - then, otherwise, .. - } => { - traverse_tree_with(then, tree_path, current_depth + 1, index_count.next(), with); - - traverse_tree_with( - otherwise, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::WrapClause { then, otherwise } => { - traverse_tree_with(then, tree_path, current_depth + 1, index_count.next(), with); - - traverse_tree_with( - otherwise, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::TupleClause { - then, otherwise, .. - } => { - traverse_tree_with(then, tree_path, current_depth + 1, index_count.next(), with); - - traverse_tree_with( - otherwise, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::Finally { pattern, then } => { - traverse_tree_with( - pattern, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - - traverse_tree_with(then, tree_path, current_depth + 1, index_count.next(), with); - } - AirExpression::If { - pattern, - then, - otherwise, - .. - } => { - traverse_tree_with( - pattern, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - - traverse_tree_with(then, tree_path, current_depth + 1, index_count.next(), with); - - traverse_tree_with( - otherwise, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::Constr { args, .. } => { - for arg in args { - traverse_tree_with(arg, tree_path, current_depth + 1, index_count.next(), with); - } - } - AirExpression::RecordUpdate { record, args, .. } => { - traverse_tree_with( - record, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - for arg in args { - traverse_tree_with(arg, tree_path, current_depth + 1, index_count.next(), with); - } - } - AirExpression::RecordAccess { record, .. } => { - traverse_tree_with( - record, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::TupleIndex { tuple, .. } => { - traverse_tree_with( - tuple, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::Trace { msg, then, .. } => { - traverse_tree_with(msg, tree_path, current_depth + 1, index_count.next(), with); - - traverse_tree_with(then, tree_path, current_depth + 1, index_count.next(), with); - } - AirExpression::FieldsEmpty { constr } => { - traverse_tree_with( - constr, - tree_path, - current_depth + 1, - index_count.next(), - with, - ); - } - AirExpression::ListEmpty { list } => { - traverse_tree_with(list, tree_path, current_depth + 1, index_count.next(), with); - } - _ => {} - }, - _ => unreachable!(), - } - - with(air_tree, tree_path); - - tree_path.pop(); -} diff --git a/crates/aiken-lang/src/gen_uplc2/tree.rs b/crates/aiken-lang/src/gen_uplc2/tree.rs index efaed17d..ce63ce05 100644 --- a/crates/aiken-lang/src/gen_uplc2/tree.rs +++ b/crates/aiken-lang/src/gen_uplc2/tree.rs @@ -1,5 +1,5 @@ use indexmap::IndexSet; -use std::sync::Arc; +use std::{borrow::BorrowMut, slice::Iter, sync::Arc}; use uplc::{builder::EXPECT_ON_LIST, builtins::DefaultFunction}; use crate::{ @@ -10,6 +10,86 @@ use crate::{ use super::air::Air; +#[derive(Clone, Debug)] +pub struct TreePath { + path: Vec<(usize, usize)>, +} + +impl TreePath { + pub fn new() -> Self { + TreePath { path: vec![] } + } + + pub fn push(&mut self, depth: usize, index: usize) { + self.path.push((depth, index)); + } + + pub fn pop(&mut self) { + self.path.pop(); + } + + pub fn current_path(&self) -> Self { + let mut path = self.path.clone(); + path.pop(); + TreePath { path } + } + + pub fn common_ancestor(&self, other: &Self) -> Self { + let mut common_ancestor = TreePath::new(); + + let mut self_iter = self.path.iter(); + let mut other_iter = other.path.iter(); + + let mut self_next = self_iter.next(); + let mut other_next = other_iter.next(); + + while self_next.is_some() && other_next.is_some() { + let self_next_level = self_next.unwrap(); + let other_next_level = other_next.unwrap(); + + if self_next_level == other_next_level { + common_ancestor.push(self_next_level.0, self_next_level.1); + } else { + break; + } + + self_next = self_iter.next(); + other_next = other_iter.next(); + } + + common_ancestor + } +} + +impl Default for TreePath { + fn default() -> Self { + Self::new() + } +} + +pub struct IndexCounter { + current_index: usize, +} + +impl IndexCounter { + pub fn new() -> Self { + IndexCounter { current_index: 0 } + } + + /// Returns the next of this [`IndexCounter`]. + pub fn next_number(&mut self) -> usize { + let current_index = self.current_index; + self.current_index += 1; + current_index + } +} + +impl Default for IndexCounter { + fn default() -> Self { + Self::new() + } +} + #[derive(Debug, Clone, PartialEq)] pub enum AirTree { Statement { @@ -144,6 +224,7 @@ pub enum AirExpression { tipo: Arc, left: Box, right: Box, + argument_tipo: Arc, }, UnOp { op: UnOp, @@ -164,6 +245,7 @@ pub enum AirExpression { tipo: Arc, subject_name: String, subject: Box, + subject_tipo: Arc, clauses: Box, }, Clause { @@ -349,12 +431,19 @@ impl AirTree { pub fn builtin(func: DefaultFunction, tipo: Arc, args: Vec) -> AirTree { AirTree::Expression(AirExpression::Builtin { func, tipo, args }) } - pub fn binop(op: BinOp, tipo: Arc, left: AirTree, right: AirTree) -> AirTree { + pub fn binop( + op: BinOp, + tipo: Arc, + left: AirTree, + right: AirTree, + argument_tipo: Arc, + ) -> AirTree { AirTree::Expression(AirExpression::BinOp { name: op, tipo, left: left.into(), right: right.into(), + argument_tipo, }) } pub fn unop(op: UnOp, arg: AirTree) -> AirTree { @@ -405,6 +494,7 @@ impl AirTree { pub fn when( subject_name: impl ToString, tipo: Arc, + subject_tipo: Arc, subject: AirTree, clauses: AirTree, ) -> AirTree { @@ -412,6 +502,7 @@ impl AirTree { tipo, subject_name: subject_name.to_string(), subject: subject.into(), + subject_tipo, clauses: clauses.into(), }) } @@ -953,10 +1044,12 @@ impl AirTree { tipo, left, right, + argument_tipo, } => { air_vec.push(Air::BinOp { name: *name, tipo: tipo.clone(), + argument_tipo: argument_tipo.clone(), }); left.create_air_vec(air_vec); right.create_air_vec(air_vec); @@ -977,11 +1070,13 @@ impl AirTree { tipo, subject_name, subject, + subject_tipo, clauses, } => { air_vec.push(Air::When { tipo: tipo.clone(), subject_name: subject_name.clone(), + subject_tipo: subject_tipo.clone(), }); subject.create_air_vec(air_vec); clauses.create_air_vec(air_vec); @@ -1135,12 +1230,12 @@ impl AirTree { } } - pub fn get_type(&self) -> Arc { + pub fn return_type(&self) -> Arc { match self { AirTree::Statement { hoisted_over: Some(hoisted_over), .. - } => hoisted_over.get_type(), + } => hoisted_over.return_type(), AirTree::Expression(e) => match e { AirExpression::Int { .. } => int(), AirExpression::String { .. } => string(), @@ -1162,7 +1257,7 @@ impl AirTree { | AirExpression::Trace { tipo, .. } => tipo.clone(), AirExpression::Void => void(), AirExpression::Var { constructor, .. } => constructor.tipo.clone(), - AirExpression::Fn { func_body, .. } => func_body.get_type(), + AirExpression::Fn { func_body, .. } => func_body.return_type(), AirExpression::UnOp { op, .. } => match op { UnOp::Not => bool(), UnOp::Negate => int(), @@ -1172,12 +1267,791 @@ impl AirTree { | AirExpression::ListClause { then, .. } | AirExpression::WrapClause { then, .. } | AirExpression::TupleClause { then, .. } - | AirExpression::Finally { then, .. } => then.get_type(), + | AirExpression::Finally { then, .. } => then.return_type(), - AirExpression::FieldsEmpty { constr } => constr.get_type(), - AirExpression::ListEmpty { list } => list.get_type(), + AirExpression::FieldsEmpty { constr } => constr.return_type(), + AirExpression::ListEmpty { list } => list.return_type(), }, _ => unreachable!(), } } + + pub fn mut_held_types(&mut self) -> Vec<&mut Arc> { + match self { + AirTree::Statement { + statement, + hoisted_over: Some(_), + } => match statement { + AirStatement::ClauseGuard { subject_tipo, .. } + | AirStatement::ListClauseGuard { subject_tipo, .. } + | AirStatement::TupleGuard { subject_tipo, .. } => vec![subject_tipo], + AirStatement::ListAccessor { tipo, .. } + | AirStatement::ListExpose { tipo, .. } + | AirStatement::TupleAccessor { tipo, .. } => vec![tipo], + _ => vec![], + }, + AirTree::Expression(e) => match e { + AirExpression::List { tipo, .. } + | AirExpression::Tuple { tipo, .. } + | AirExpression::Call { tipo, .. } + | AirExpression::Builtin { tipo, .. } + | AirExpression::UnWrapData { tipo, .. } + | AirExpression::WrapData { tipo, .. } + | AirExpression::If { tipo, .. } + | AirExpression::RecordUpdate { tipo, .. } + | AirExpression::RecordAccess { tipo, .. } + | AirExpression::Constr { tipo, .. } + | AirExpression::TupleIndex { tipo, .. } + | AirExpression::ErrorTerm { tipo } + | AirExpression::Trace { tipo, .. } => vec![tipo], + AirExpression::Var { constructor, .. } => { + vec![constructor.tipo.borrow_mut()] + } + AirExpression::BinOp { + tipo, + argument_tipo, + .. + } => { + vec![tipo, argument_tipo] + } + AirExpression::When { + tipo, subject_tipo, .. + } => vec![tipo, subject_tipo], + AirExpression::Clause { subject_tipo, .. } + | AirExpression::ListClause { subject_tipo, .. } + | AirExpression::TupleClause { subject_tipo, .. } => vec![subject_tipo], + _ => { + vec![] + } + }, + _ => unreachable!("FOUND UNHOISTED STATEMENT"), + } + } + + pub fn traverse_tree_with(&mut self, with: &mut impl FnMut(&mut AirTree, &TreePath)) { + let mut tree_path = TreePath::new(); + self.do_traverse_tree_with(&mut tree_path, 0, 0, with); + } + + fn do_traverse_tree_with( + &mut self, + tree_path: &mut TreePath, + current_depth: usize, + depth_index: usize, + with: &mut impl FnMut(&mut AirTree, &TreePath), + ) { + let mut index_count = IndexCounter::new(); + tree_path.push(current_depth, depth_index); + match self { + AirTree::Statement { + statement, + hoisted_over: Some(hoisted_over), + } => { + match statement { + AirStatement::Let { value, .. } => { + value.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirStatement::DefineFunc { func_body, .. } => { + func_body.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirStatement::AssertConstr { constr, .. } => { + constr.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirStatement::AssertBool { value, .. } => { + value.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirStatement::ClauseGuard { pattern, .. } => { + pattern.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirStatement::ListClauseGuard { .. } => {} + AirStatement::TupleGuard { .. } => {} + AirStatement::FieldsExpose { record, .. } => { + record.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirStatement::ListAccessor { list, .. } => { + list.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirStatement::ListExpose { .. } => {} + AirStatement::TupleAccessor { tuple, .. } => { + tuple.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirStatement::NoOp => {} + }; + + hoisted_over.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirTree::Expression(e) => match e { + AirExpression::List { items, .. } => { + for item in items { + item.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + } + AirExpression::Tuple { items, .. } => { + for item in items { + item.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + } + AirExpression::Call { func, args, .. } => { + func.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + for arg in args { + arg.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + } + AirExpression::Fn { func_body, .. } => { + func_body.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::Builtin { args, .. } => { + for arg in args { + arg.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + } + AirExpression::BinOp { left, right, .. } => { + left.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + right.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::UnOp { arg, .. } => { + arg.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::UnWrapData { value, .. } => { + value.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::WrapData { value, .. } => { + value.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::When { + subject, clauses, .. + } => { + subject.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + clauses.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::Clause { + pattern, + then, + otherwise, + .. + } => { + pattern.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + then.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + otherwise.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::ListClause { + then, otherwise, .. + } => { + then.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + otherwise.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::WrapClause { then, otherwise } => { + then.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + otherwise.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::TupleClause { + then, otherwise, .. + } => { + then.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + otherwise.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::Finally { pattern, then } => { + pattern.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + then.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::If { + pattern, + then, + otherwise, + .. + } => { + pattern.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + then.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + otherwise.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::Constr { args, .. } => { + for arg in args { + arg.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + } + AirExpression::RecordUpdate { record, args, .. } => { + record.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + for arg in args { + arg.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + } + AirExpression::RecordAccess { record, .. } => { + record.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::TupleIndex { tuple, .. } => { + tuple.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::Trace { msg, then, .. } => { + msg.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + + then.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::FieldsEmpty { constr } => { + constr.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + AirExpression::ListEmpty { list } => { + list.do_traverse_tree_with( + tree_path, + current_depth + 1, + index_count.next_number(), + with, + ); + } + _ => {} + }, + _ => unreachable!(), + } + + with(self, tree_path); + + tree_path.pop(); + } + + pub fn find_air_tree_node<'a>(&'a mut self, tree_path: &TreePath) -> &'a mut AirTree { + let mut path_iter = tree_path.path.iter(); + self.do_find_air_tree_node(&mut path_iter) + } + + fn do_find_air_tree_node<'a>( + &'a mut self, + tree_path_iter: &mut Iter<(usize, usize)>, + ) -> &'a mut AirTree { + if let Some((_depth, index)) = tree_path_iter.next() { + let mut children_nodes = vec![]; + match self { + AirTree::Statement { + statement, + hoisted_over: Some(hoisted_over), + } => match statement { + AirStatement::Let { value, .. } => { + if *index == 0 { + value.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::DefineFunc { func_body, .. } => { + if *index == 0 { + func_body.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::AssertConstr { constr, .. } => { + if *index == 0 { + constr.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::AssertBool { value, .. } => { + if *index == 0 { + value.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::ClauseGuard { pattern, .. } => { + if *index == 0 { + pattern.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::ListClauseGuard { .. } => { + if *index == 0 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::TupleGuard { .. } => { + if *index == 0 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::FieldsExpose { record, .. } => { + if *index == 0 { + record.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::ListAccessor { list, .. } => { + if *index == 0 { + list.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::ListExpose { .. } => { + if *index == 0 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::TupleAccessor { tuple, .. } => { + if *index == 0 { + tuple.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirStatement::NoOp => { + if *index == 0 { + hoisted_over.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + }, + AirTree::Expression(e) => match e { + AirExpression::List { items, .. } + | AirExpression::Tuple { items, .. } + | AirExpression::Builtin { args: items, .. } => { + let item = items.get_mut(*index).unwrap_or_else(|| { + panic!("Tree Path index outside tree children nodes") + }); + item.do_find_air_tree_node(tree_path_iter) + } + AirExpression::Call { func, args, .. } => { + children_nodes.push(func.as_mut()); + children_nodes.extend(args.iter_mut()); + + let item = children_nodes.swap_remove(*index); + + item.do_find_air_tree_node(tree_path_iter) + } + AirExpression::Fn { func_body, .. } => { + if *index == 0 { + func_body.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::BinOp { left, right, .. } => { + if *index == 0 { + left.as_mut() + } else if *index == 1 { + right.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::UnOp { arg, .. } => { + if *index == 0 { + arg.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::UnWrapData { value, .. } => { + if *index == 0 { + value.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::WrapData { value, .. } => { + if *index == 0 { + value.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::When { + subject, clauses, .. + } => { + if *index == 0 { + subject.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + clauses.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::Clause { + pattern, + then, + otherwise, + .. + } => { + if *index == 0 { + pattern.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + then.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 2 { + otherwise.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::ListClause { + then, otherwise, .. + } => { + if *index == 0 { + then.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + otherwise.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::WrapClause { then, otherwise } => { + if *index == 0 { + then.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + otherwise.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::TupleClause { + then, otherwise, .. + } => { + if *index == 0 { + then.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + otherwise.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::Finally { pattern, then } => { + if *index == 0 { + pattern.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + then.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::If { + pattern, + then, + otherwise, + .. + } => { + if *index == 0 { + pattern.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + then.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 2 { + otherwise.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::Constr { args, .. } => { + let item = args.get_mut(*index).unwrap_or_else(|| { + panic!("Tree Path index outside tree children nodes") + }); + item.do_find_air_tree_node(tree_path_iter) + } + AirExpression::RecordUpdate { record, args, .. } => { + children_nodes.push(record.as_mut()); + children_nodes.extend(args.iter_mut()); + + let item = children_nodes.swap_remove(*index); + + item.do_find_air_tree_node(tree_path_iter) + } + AirExpression::RecordAccess { record, .. } => { + if *index == 0 { + record.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::TupleIndex { tuple, .. } => { + if *index == 0 { + tuple.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + + AirExpression::Trace { msg, then, .. } => { + if *index == 0 { + msg.as_mut().do_find_air_tree_node(tree_path_iter) + } else if *index == 1 { + then.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::FieldsEmpty { constr } => { + if *index == 0 { + constr.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + AirExpression::ListEmpty { list } => { + if *index == 0 { + list.as_mut().do_find_air_tree_node(tree_path_iter) + } else { + panic!("Tree Path index outside tree children nodes") + } + } + _ => unreachable!( + "A tree node with no children was encountered with a longer tree path." + ), + }, + _ => unreachable!(), + } + } else { + self + } + } }