From 61184fbb86be4cf8c388f2f4c57047663508588e Mon Sep 17 00:00:00 2001 From: microproofs Date: Tue, 22 Oct 2024 17:04:29 -0400 Subject: [PATCH] Use Stick breaking set to track previously encountered values for each DecisionTree switch case --- crates/aiken-lang/src/gen_uplc.rs | 84 +++++- .../aiken-lang/src/gen_uplc/decision_tree.rs | 78 ++++-- .../src/gen_uplc/stick_break_set.rs | 243 ++++++++++++++++++ 3 files changed, 370 insertions(+), 35 deletions(-) create mode 100644 crates/aiken-lang/src/gen_uplc/stick_break_set.rs diff --git a/crates/aiken-lang/src/gen_uplc.rs b/crates/aiken-lang/src/gen_uplc.rs index 12c80c71..9b821b7f 100644 --- a/crates/aiken-lang/src/gen_uplc.rs +++ b/crates/aiken-lang/src/gen_uplc.rs @@ -2,14 +2,15 @@ pub mod air; pub mod builder; pub mod decision_tree; pub mod interner; +pub mod stick_break_set; pub mod tree; use self::{ air::Air, builder::{ cast_validator_args, convert_type_to_data, extract_constant, modify_cyclic_calls, - modify_self_calls, rearrange_list_clauses, AssignmentProperties, ClauseProperties, - CodeGenSpecialFuncs, CycleFunctionNames, HoistableFunction, Variant, + modify_self_calls, AssignmentProperties, ClauseProperties, CodeGenSpecialFuncs, + CycleFunctionNames, HoistableFunction, Variant, }, tree::{AirTree, TreePath}, }; @@ -44,12 +45,13 @@ use builder::{ introduce_name, introduce_pattern, pop_pattern, softcast_data_to_type_otherwise, unknown_data_to_type, DISCARDED, }; -use decision_tree::{get_tipo_by_path, name_from_path, TreeGen}; +use decision_tree::{get_tipo_by_path, TreeGen}; use indexmap::{IndexMap, IndexSet}; use interner::AirInterner; use itertools::Itertools; use petgraph::{algo, Graph}; use std::{collections::HashMap, rc::Rc}; +use stick_break_set::{Builtin, Builtins, TreeSet}; use tree::Fields; use uplc::{ ast::{Constant as UplcConstant, Name, NamedDeBruijn, Program, Term, Type as UplcType}, @@ -624,7 +626,9 @@ impl<'a> CodeGenerator<'a> { let tree = tree_gen.build_tree(&subject_name_interned, &subject.tipo(), clauses); - let clauses = self.handle_decision_tree(&constr_var_interned, tree); + let stick_set = TreeSet::new(); + + let clauses = self.handle_decision_tree(tree, stick_set); self.interner.pop_text(constr_var); self.interner.pop_text(subject_name); @@ -2456,9 +2460,9 @@ impl<'a> CodeGenerator<'a> { } fn handle_decision_tree( - &self, - constr_name: &String, + &mut self, tree: decision_tree::DecisionTree<'_>, + mut stick_set: TreeSet, ) -> AirTree { match tree { decision_tree::DecisionTree::Switch { @@ -2468,11 +2472,34 @@ impl<'a> CodeGenerator<'a> { mut cases, default, } => { - let current_tipo = get_tipo_by_path(subject_tipo, &path); + let current_tipo = get_tipo_by_path(subject_tipo.clone(), &path); + + let builtins_path = Builtins::new_from_path(subject_tipo.clone(), path); + + let mut prev_builtins = builtins_path.clone(); + + let builtins_to_add = stick_set.diff_union_builtins(builtins_path.clone()); + + builtins_to_add.vec.iter().for_each(|_| { + prev_builtins.vec.pop(); + }); + + let prev_subject_name = if prev_builtins.is_empty() { + subject_name.clone() + } else { + format!("{}_{}", subject_name, prev_builtins.to_string()) + }; + let prev_tipo = prev_builtins.vec.last().unwrap().tipo(); + + let current_subject_name = if builtins_path.is_empty() { + subject_name + } else { + format!("{}_{}", subject_name, builtins_path.to_string()) + }; let data_type = lookup_data_type_by_tipo(&self.data_types, ¤t_tipo); - let needs_default = if let Some(data_type) = data_type { + let needs_default = if let Some(data_type) = &data_type { data_type.constructors.len() != cases.len() } else { true @@ -2484,9 +2511,46 @@ impl<'a> CodeGenerator<'a> { cases.pop().unwrap().1 }; - let last_then = self.handle_decision_tree(constr_name, last_then); + let last_then = AirTree::anon_func( + vec![], + self.handle_decision_tree(last_then, stick_set.clone()), + true, + ); - todo!() + let test_subject_name = if data_type.is_some() { + format!("{}_index", current_subject_name.clone(),) + } else { + current_subject_name.clone() + }; + + let clauses = cases.into_iter().rfold(last_then, |acc, (case, then)| { + let case_air = AirTree::anon_func( + vec![], + self.handle_decision_tree(then, stick_set.clone()), + true, + ); + + AirTree::clause( + test_subject_name.clone(), + case.get_air_pattern(), + current_tipo.clone(), + case_air, + acc, + false, + ) + }); + + let y = AirTree::when( + test_subject_name, + Type::void(), + current_tipo.clone(), + AirTree::local_var(current_subject_name, current_tipo.clone()), + clauses, + ); + + let x = builtins_to_add.to_air(prev_subject_name, prev_tipo, y); + + x } decision_tree::DecisionTree::ListSwitch { subject_name, diff --git a/crates/aiken-lang/src/gen_uplc/decision_tree.rs b/crates/aiken-lang/src/gen_uplc/decision_tree.rs index 071f547d..ee36184f 100644 --- a/crates/aiken-lang/src/gen_uplc/decision_tree.rs +++ b/crates/aiken-lang/src/gen_uplc/decision_tree.rs @@ -10,7 +10,7 @@ use crate::{ expr::{lookup_data_type_by_tipo, Type, TypedExpr}, }; -use super::interner::AirInterner; +use super::{interner::AirInterner, tree::AirTree}; const PAIR_NEW_COLUMNS: usize = 2; @@ -31,18 +31,40 @@ pub enum Path { ListTail(usize), } -impl Display for Path { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { +impl ToString for Path { + fn to_string(&self) -> String { match self { - Path::Pair(i) => write!(f, "Pair({})", i), - Path::Tuple(i) => write!(f, "Tuple({})", i), - Path::Constr(_, i) => write!(f, "Constr({})", i), - Path::List(i) => write!(f, "List({})", i), - Path::ListTail(i) => write!(f, "ListTail({})", i), + Path::Pair(i) => { + format!("pair_{}", i) + } + Path::Tuple(i) => { + format!("tuple_{}", i) + } + Path::Constr(_, i) => { + format!("constr_{}", i) + } + Path::List(i) => { + format!("list_{}", i) + } + Path::ListTail(i) => { + format!("listtail_{}", i) + } } } } +// impl Display for Path { +// fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { +// match self { +// Path::Pair(i) => write!(f, "Pair({})", i), +// Path::Tuple(i) => write!(f, "Tuple({})", i), +// Path::Constr(_, i) => write!(f, "Constr({})", i), +// Path::List(i) => write!(f, "List({})", i), +// Path::ListTail(i) => write!(f, "ListTail({})", i), +// } +// } +// } + impl PartialEq for Path { fn eq(&self, other: &Self) -> bool { match (self, other) { @@ -92,6 +114,19 @@ pub enum CaseTest { Wild, } +impl CaseTest { + pub fn get_air_pattern(&self) -> AirTree { + match self { + CaseTest::Constr(i) => AirTree::int(i), + CaseTest::Int(i) => AirTree::int(i), + CaseTest::Bytes(vec) => AirTree::byte_array(vec.clone()), + CaseTest::List(_) => unreachable!(), + CaseTest::ListWithTail(_) => unreachable!(), + CaseTest::Wild => unreachable!(), + } + } +} + impl Display for CaseTest { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { @@ -231,7 +266,10 @@ impl<'a> DecisionTree<'a> { .append( path.iter() .fold(RcDoc::line().append(RcDoc::text("path(")), |acc, p| { - acc.append(RcDoc::line().append(RcDoc::text(format!("{}", p)).nest(4))) + acc.append( + RcDoc::line() + .append(RcDoc::text(format!("{}", p.to_string())).nest(4)), + ) }) .append(RcDoc::line()) .append(RcDoc::text(")")) @@ -277,7 +315,10 @@ impl<'a> DecisionTree<'a> { .append( path.iter() .fold(RcDoc::line().append(RcDoc::text("path(")), |acc, p| { - acc.append(RcDoc::line().append(RcDoc::text(format!("{}", p)).nest(4))) + acc.append( + RcDoc::line() + .append(RcDoc::text(format!("{}", p.to_string())).nest(4)), + ) }) .append(RcDoc::line()) .append(RcDoc::text(")")) @@ -1228,10 +1269,6 @@ fn highest_occurrence(matrix: &PatternMatrix, column_length: usize) -> Option) -> String { - todo!() -} - #[cfg(test)] mod tester { use std::collections::HashMap; @@ -1244,7 +1281,7 @@ mod tester { UntypedModule, }, builtins, - expr::{Type, TypedExpr}, + expr::TypedExpr, gen_uplc::{decision_tree::TreeGen, interner::AirInterner}, parser, tipo::error::{Error, Warning}, @@ -1426,16 +1463,7 @@ mod tester { let tree_gen = TreeGen::new(&mut air_interner, &data_types, &pattern); - let tree = tree_gen.build_tree( - &"subject".to_string(), - &Type::tuple(vec![ - Type::int(), - Type::int(), - Type::byte_array(), - Type::list(Type::int()), - ]), - clauses, - ); + let tree = tree_gen.build_tree(&"subject".to_string(), &subject.tipo(), clauses); println!("{}", tree); } diff --git a/crates/aiken-lang/src/gen_uplc/stick_break_set.rs b/crates/aiken-lang/src/gen_uplc/stick_break_set.rs new file mode 100644 index 00000000..f5163449 --- /dev/null +++ b/crates/aiken-lang/src/gen_uplc/stick_break_set.rs @@ -0,0 +1,243 @@ +use std::rc::Rc; + +use itertools::Itertools; +use uplc::builtins::DefaultFunction; + +use crate::expr::Type; + +use super::{ + decision_tree::{get_tipo_by_path, Path}, + tree::AirTree, +}; + +#[derive(Clone)] +pub enum Builtin { + HeadList(Rc), + TailList, + UnConstr, + FstPair(Rc), + SndPair(Rc), +} + +impl PartialEq for Builtin { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Builtin::HeadList(_), Builtin::HeadList(_)) => true, + (Builtin::TailList, Builtin::TailList) => true, + (Builtin::UnConstr, Builtin::UnConstr) => true, + (Builtin::SndPair(_), Builtin::SndPair(_)) => true, + _ => false, + } + } +} + +impl Eq for Builtin {} + +impl Builtin { + fn to_air_call(self, arg: AirTree) -> AirTree { + match self { + Builtin::HeadList(t) => AirTree::builtin(DefaultFunction::HeadList, t, vec![arg]), + Builtin::TailList => AirTree::builtin( + DefaultFunction::TailList, + Type::list(Type::data()), + vec![arg], + ), + Builtin::UnConstr => AirTree::builtin( + DefaultFunction::TailList, + Type::pair(Type::int(), Type::list(Type::data())), + vec![arg], + ), + Builtin::FstPair(t) => AirTree::builtin(DefaultFunction::FstPair, t, vec![arg]), + Builtin::SndPair(t) => AirTree::builtin(DefaultFunction::SndPair, t, vec![arg]), + } + } + + pub fn tipo(&self) -> Rc { + match self { + Builtin::HeadList(t) => t.clone(), + Builtin::TailList => Type::list(Type::data()), + + Builtin::UnConstr => Type::pair(Type::int(), Type::list(Type::data())), + + Builtin::FstPair(t) => t.clone(), + Builtin::SndPair(t) => t.clone(), + } + } +} + +impl ToString for Builtin { + fn to_string(&self) -> String { + match self { + Builtin::HeadList(_) => "head".to_string(), + Builtin::TailList => "tail".to_string(), + Builtin::UnConstr => "unconstr".to_string(), + Builtin::FstPair(_) => "fst".to_string(), + Builtin::SndPair(_) => "snd".to_string(), + } + } +} + +#[derive(PartialEq, Eq, Clone)] +pub struct Builtins { + pub vec: Vec, +} + +impl Builtins { + pub fn new() -> Self { + Builtins { vec: vec![] } + } + + pub fn new_from_path(subject_tipo: Rc, path: Vec) -> Self { + Self { + vec: path + .into_iter() + .fold((vec![], vec![]), |(mut builtins, mut rebuilt_path), i| { + rebuilt_path.push(i.clone()); + match i { + Path::Pair(i) => { + if i == 0 { + builtins.push(Builtin::HeadList(get_tipo_by_path( + subject_tipo.clone(), + &rebuilt_path, + ))); + } else if i == 1 { + builtins.push(Builtin::SndPair(get_tipo_by_path( + subject_tipo.clone(), + &rebuilt_path, + ))); + } else { + unreachable!() + } + + (builtins, rebuilt_path) + } + Path::List(i) | Path::Tuple(i) => { + for _ in 0..i { + builtins.push(Builtin::TailList); + } + + builtins.push(Builtin::HeadList(get_tipo_by_path( + subject_tipo.clone(), + &rebuilt_path, + ))); + + (builtins, rebuilt_path) + } + Path::Constr(_rc, i) => { + builtins.extend([ + Builtin::UnConstr, + Builtin::SndPair(Type::list(Type::data())), + ]); + + for _ in 0..i { + builtins.push(Builtin::TailList); + } + + builtins.push(Builtin::HeadList(get_tipo_by_path( + subject_tipo.clone(), + &rebuilt_path, + ))); + + (builtins, rebuilt_path) + } + + Path::ListTail(i) => { + for _ in 0..i { + builtins.push(Builtin::TailList); + } + + (builtins, rebuilt_path) + } + } + }) + .0, + } + } + + pub fn is_empty(&self) -> bool { + self.vec.is_empty() + } + + pub fn to_air(self, prev_name: String, subject_tipo: Rc, then: AirTree) -> AirTree { + let (_, _, name_builtins) = self.vec.into_iter().fold( + (prev_name, subject_tipo, vec![]), + |(prev_name, prev_tipo, mut acc), item| { + let next_name = format!("{}_{}", prev_name, item.to_string()); + let next_tipo = item.tipo(); + + acc.push((prev_name, prev_tipo, next_name.clone(), item)); + + (next_name, next_tipo, acc) + }, + ); + + name_builtins + .into_iter() + .rfold(then, |then, (prev_name, prev_tipo, next_name, builtin)| { + AirTree::let_assignment( + next_name, + builtin.to_air_call(AirTree::local_var(prev_name, prev_tipo)), + then, + ) + }) + } +} + +impl ToString for Builtins { + fn to_string(&self) -> String { + self.vec.iter().map(|i| i.to_string()).join("_") + } +} + +#[derive(Clone)] +pub struct TreeSet { + children: Vec, +} + +#[derive(Clone)] +pub struct TreeNode { + node: Builtin, + children: Vec, +} + +impl TreeNode {} + +impl TreeSet { + pub fn new() -> Self { + TreeSet { children: vec![] } + } + + pub fn new_from_builtins(builtins: Builtins) -> Self { + TreeSet { + children: builtins + .vec + .into_iter() + .map(|item| TreeNode { + node: item, + children: vec![], + }) + .rev() + .reduce(|prev, mut current| { + current.children.push(prev); + current + }) + .into_iter() + .collect_vec(), + } + } + + pub fn diff_union_builtins(&mut self, builtins: Builtins) -> Builtins { + if let Some((first, rest)) = builtins.vec.split_first() { + if self.children.iter().any(|item| first == &item.node) { + todo!() + } else { + self.children + .extend(TreeSet::new_from_builtins(builtins.clone()).children); + + builtins + } + } else { + builtins + } + } +}