Use Stick breaking set to track previously encountered values for each DecisionTree switch case

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,
+        modify_self_calls, AssignmentProperties, ClauseProperties, CodeGenSpecialFuncs,
-        CodeGenSpecialFuncs, CycleFunctionNames, HoistableFunction, Variant,
+        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,
+        &mut self,
-        constr_name: &String,
         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, &current_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,

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,17 +31,39 @@ pub enum Path {
     ListTail(usize),
 }
 
-impl Display for Path {
+impl ToString for Path {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    fn to_string(&self) -> String {
         match self {
-            Path::Pair(i) => write!(f, "Pair({})", i),
+            Path::Pair(i) => {
-            Path::Tuple(i) => write!(f, "Tuple({})", i),
+                format!("pair_{}", i)
-            Path::Constr(_, i) => write!(f, "Constr({})", i),
+            }
-            Path::List(i) => write!(f, "List({})", i),
+            Path::Tuple(i) => {
-            Path::ListTail(i) => write!(f, "ListTail({})", 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 {
@@ -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<us
     }
 }
 
-pub fn name_from_path(subject_name: String, path: Vec<Path>) -> 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(
+        let tree = tree_gen.build_tree(&"subject".to_string(), &subject.tipo(), clauses);
-            &"subject".to_string(),
-            &Type::tuple(vec![
-                Type::int(),
-                Type::int(),
-                Type::byte_array(),
-                Type::list(Type::int()),
-            ]),
-            clauses,
-        );
 
         println!("{}", tree);
     }

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<Type>),
+    TailList,
+    UnConstr,
+    FstPair(Rc<Type>),
+    SndPair(Rc<Type>),
+}
+
+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<Type> {
+        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<Builtin>,
+}
+
+impl Builtins {
+    pub fn new() -> Self {
+        Builtins { vec: vec![] }
+    }
+
+    pub fn new_from_path(subject_tipo: Rc<Type>, path: Vec<Path>) -> 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<Type>, 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<TreeNode>,
+}
+
+#[derive(Clone)]
+pub struct TreeNode {
+    node: Builtin,
+    children: Vec<TreeNode>,
+}
+
+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
+        }
+    }
+}