Finish up decision tree and rework it a bit to closely follow how the paper handles wild card patterns

Cargo.lock

@@ -99,6 +99,7 @@ dependencies = [
  "pallas-primitives",
  "patricia_tree",
  "petgraph",
+ "pretty 0.12.3",
  "pretty_assertions",
  "serde",
  "strum",
@@ -2241,6 +2242,17 @@ dependencies = [
  "unicode-segmentation",
 ]
 
+[[package]]
+name = "pretty"
+version = "0.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b55c4d17d994b637e2f4daf6e5dc5d660d209d5642377d675d7a1c3ab69fa579"
+dependencies = [
+ "arrayvec",
+ "typed-arena",
+ "unicode-width",
+]
+
 [[package]]
 name = "pretty_assertions"
 version = "1.4.1"
@@ -3254,7 +3266,7 @@ dependencies = [
  "pallas-primitives",
  "pallas-traverse",
  "peg",
- "pretty",
+ "pretty 0.11.3",
  "pretty_assertions",
  "secp256k1",
  "serde",

crates/aiken-lang/Cargo.toml

@@ -28,6 +28,7 @@ owo-colors = { version = "3.5.0", features = ["supports-colors"] }
 pallas-primitives.workspace = true
 patricia_tree = "0.8.0"
 petgraph = "0.6.3"
+pretty = "0.12.3"
 serde = { version = "1.0.197", features = ["derive", "rc"] }
 strum = "0.24.1"
 thiserror = "1.0.39"

crates/aiken-lang/src/gen_uplc/decision_tree.rs

@@ -1,4 +1,6 @@
-use std::rc::Rc;
+use core::fmt;
+use pretty::RcDoc;
+use std::{fmt::Display, rc::Rc};
 
 use indexmap::IndexMap;
 use itertools::{Itertools, Position};
@@ -29,6 +31,18 @@ pub enum Path {
     ListTail(usize),
 }
 
+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) {
@@ -42,6 +56,8 @@ impl PartialEq for Path {
     }
 }
 
+impl Eq for Path {}
+
 #[derive(Clone, Debug)]
 pub struct Assigned {
     path: Vec<Path>,
@@ -58,7 +74,7 @@ struct RowItem<'a> {
 struct Row<'a> {
     assigns: Vec<Assigned>,
     columns: Vec<RowItem<'a>>,
-    then: &'a TypedExpr,
+    then: String,
 }
 
 #[derive(Clone, Debug)]
@@ -76,6 +92,19 @@ pub enum CaseTest {
     Wild,
 }
 
+impl Display for CaseTest {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CaseTest::Constr(i) => write!(f, "Constr({})", i),
+            CaseTest::Int(i) => write!(f, "Int({})", i),
+            CaseTest::Bytes(vec) => write!(f, "Bytes({:?})", vec),
+            CaseTest::List(i) => write!(f, "List({})", i),
+            CaseTest::ListWithTail(i) => write!(f, "ListWithTail({})", i),
+            CaseTest::Wild => write!(f, "Wild"),
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub enum DecisionTree<'a> {
     Switch {
@@ -83,7 +112,7 @@ pub enum DecisionTree<'a> {
         subject_tipo: Rc<Type>,
         path: Vec<Path>,
         cases: Vec<(CaseTest, DecisionTree<'a>)>,
-        default: Box<DecisionTree<'a>>,
+        default: Option<Box<DecisionTree<'a>>>,
     },
     ListSwitch {
         subject_name: String,
@@ -93,40 +122,217 @@ pub enum DecisionTree<'a> {
         tail_cases: Vec<(CaseTest, DecisionTree<'a>)>,
         default: Option<Box<DecisionTree<'a>>>,
     },
-    Leaf(Vec<Assigned>, &'a TypedExpr),
+    HoistedLeaf(String, Vec<Assigned>),
-    HoistedLeaf(String),
+    HoistThen {
-    HoistThen(String, Box<DecisionTree<'a>>, Box<DecisionTree<'a>>),
+        name: String,
+        assigns: Vec<Assigned>,
+        pattern: Box<DecisionTree<'a>>,
+        then: &'a TypedExpr,
+    },
+}
+
+impl<'a> DecisionTree<'a> {
+    pub fn to_pretty(&self) -> String {
+        let mut w = Vec::new();
+
+        self.to_doc().render(80, &mut w).unwrap();
+
+        String::from_utf8(w)
+            .unwrap()
+            .lines()
+            // This is a hack to deal with blank newlines
+            // that end up with a bunch of useless whitespace
+            // because of the nesting
+            .map(|l| {
+                if l.chars().all(|c| c.is_whitespace()) {
+                    "".to_string()
+                } else {
+                    l.to_string()
+                }
+            })
+            .collect::<Vec<_>>()
+            .join("\n")
+    }
+
+    fn to_doc(&self) -> RcDoc<()> {
+        match self {
+            DecisionTree::Switch {
+                path,
+                cases,
+                default,
+                ..
+            } => RcDoc::text("Switch(")
+                .append(RcDoc::line())
+                .append(
+                    path.iter()
+                        .fold(RcDoc::text("path("), |acc, p| {
+                            acc.append(RcDoc::line())
+                                .append(RcDoc::text(format!("{}", p)))
+                        })
+                        .append(RcDoc::line_())
+                        .nest(2)
+                        .append(RcDoc::text(")")),
+                )
+                .append(RcDoc::line())
+                .append(
+                    cases
+                        .iter()
+                        .fold(RcDoc::text("cases("), |acc, (con, tree)| {
+                            acc.append(RcDoc::line())
+                                .append(format!("({}): ", con))
+                                .append(RcDoc::line())
+                                .append(tree.to_doc().nest(2))
+                        })
+                        .append(RcDoc::line_())
+                        .nest(2)
+                        .append(RcDoc::text(")")),
+                )
+                .append(RcDoc::line())
+                .append(
+                    RcDoc::text("default : ")
+                        .append(RcDoc::line())
+                        .append(
+                            default
+                                .as_ref()
+                                .map(|i| i.to_doc())
+                                .unwrap_or(RcDoc::text("None")),
+                        )
+                        .nest(2),
+                )
+                .append(RcDoc::line_())
+                .append(RcDoc::text(")")),
+            DecisionTree::ListSwitch {
+                path,
+                cases,
+                tail_cases,
+                default,
+                ..
+            } => RcDoc::text("ListSwitch(")
+                .append(
+                    path.iter()
+                        .fold(RcDoc::text("path("), |acc, p| {
+                            acc.append(RcDoc::line())
+                                .append(RcDoc::text(format!("{}", p)))
+                        })
+                        .append(RcDoc::line_())
+                        .nest(2)
+                        .append(RcDoc::text(")")),
+                )
+                .append(
+                    cases
+                        .iter()
+                        .fold(RcDoc::text("cases("), |acc, (con, tree)| {
+                            acc.append(RcDoc::line())
+                                .append(format!("({}): ", con))
+                                .append(RcDoc::line())
+                                .append(tree.to_doc().nest(2))
+                        })
+                        .append(RcDoc::line_())
+                        .nest(2)
+                        .append(RcDoc::text(")")),
+                )
+                .append(
+                    tail_cases
+                        .iter()
+                        .fold(RcDoc::text("tail cases("), |acc, (con, tree)| {
+                            acc.append(RcDoc::line())
+                                .append(format!("({}): ", con))
+                                .append(RcDoc::line())
+                                .append(tree.to_doc().nest(2))
+                        })
+                        .append(RcDoc::line_())
+                        .nest(2)
+                        .append(RcDoc::text(")")),
+                )
+                .append(
+                    RcDoc::text("default : ")
+                        .append(RcDoc::line())
+                        .append(
+                            default
+                                .as_ref()
+                                .map(|i| i.to_doc())
+                                .unwrap_or(RcDoc::text("None")),
+                        )
+                        .nest(2),
+                )
+                .append(RcDoc::line_())
+                .append(RcDoc::text(")")),
+            DecisionTree::HoistedLeaf(name, _) => RcDoc::text(format!("Leaf({})", name)),
+            DecisionTree::HoistThen { .. } => todo!(),
+        }
+    }
+}
+
+impl<'a> Display for DecisionTree<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.to_pretty())
+    }
 }
 
 pub struct TreeGen<'a, 'b> {
     interner: &'b mut AirInterner,
     data_types: &'b IndexMap<&'a DataTypeKey, &'a TypedDataType>,
+    wild_card_pattern: RowItem<'a>,
 }
 
 impl<'a, 'b> TreeGen<'a, 'b> {
+    pub fn new(
+        interner: &'b mut AirInterner,
+        data_types: &'b IndexMap<&'a DataTypeKey, &'a TypedDataType>,
+        wild_card_pattern: &'a TypedPattern,
+    ) -> Self {
+        TreeGen {
+            interner,
+            data_types,
+            wild_card_pattern: RowItem {
+                path: vec![],
+                pattern: wild_card_pattern,
+            },
+        }
+    }
+
     pub fn build_tree(
         mut self,
         subject_name: &String,
         subject_tipo: &Rc<Type>,
         clauses: &'a [TypedClause],
     ) -> DecisionTree<'a> {
+        let mut clause_then_map = IndexMap::new();
+
         let rows = clauses
             .iter()
-            .map(|clause| {
+            .enumerate()
+            .map(|(index, clause)| {
                 let (assign, row_items) =
                     self.map_pattern_to_row(&clause.pattern, subject_tipo, vec![]);
 
-                Row {
+                self.interner.intern(format!("__clause_then_{}", index));
+                let clause_then_name = self
+                    .interner
+                    .lookup_interned(&format!("__clause_then_{}", index));
+
+                clause_then_map.insert(clause_then_name.clone(), (vec![], &clause.then));
+
+                let row = Row {
                     assigns: assign.into_iter().collect_vec(),
                     columns: row_items,
-                    then: &clause.then,
+                    then: clause_then_name,
-                }
+                };
+
+                self.interner.pop_text(format!("__clause_then_{}", index));
+
+                row
             })
             .collect_vec();
 
         let tree_gen = &mut self;
 
-        tree_gen.do_build_tree(subject_name, subject_tipo, PatternMatrix { rows }, None)
+        tree_gen.do_build_tree(
+            subject_name,
+            subject_tipo,
+            PatternMatrix { rows },
+            &mut clause_then_map,
+        )
     }
 
     fn do_build_tree(
@@ -134,21 +340,41 @@ impl<'a, 'b> TreeGen<'a, 'b> {
         subject_name: &String,
         subject_tipo: &Rc<Type>,
         matrix: PatternMatrix<'a>,
-        fallback_option: Option<DecisionTree<'a>>,
+        then_map: &mut IndexMap<String, (Vec<Assigned>, &'a TypedExpr)>,
     ) -> DecisionTree<'a> {
         let column_length = matrix.rows[0].columns.len();
-
+        // First step make sure all rows have same number of columns
+        // or something went wrong
         assert!(matrix
             .rows
             .iter()
             .all(|row| { row.columns.len() == column_length }));
 
         let occurrence_col = highest_occurrence(&matrix, column_length);
+        // Find which column has the most important pattern
+
+        let Some(occurrence_col) = occurrence_col else {
+            // No more patterns to match on so we grab the first default row and return that
+            let mut fallback = matrix.rows;
+
+            let row = fallback.swap_remove(0);
+
+            let Some((assigns, _)) = then_map.get_mut(&row.then) else {
+                unreachable!()
+            };
+
+            if assigns.is_empty() {
+                *assigns = row.assigns.clone();
+            }
+
+            return DecisionTree::HoistedLeaf(row.then, row.assigns);
+        };
 
         let mut longest_elems_no_tail = None;
         let mut longest_elems_with_tail = None;
         let mut has_list_pattern = false;
 
+        // List patterns are special so we need more information on length
         matrix.rows.iter().for_each(|item| {
             let col = &item.columns[occurrence_col];
 
@@ -194,19 +420,20 @@ impl<'a, 'b> TreeGen<'a, 'b> {
 
         let specialized_tipo = get_tipo_by_path(subject_tipo.clone(), &path);
 
-        let mut row_iter = matrix.rows.into_iter().peekable();
+        let (default_matrix, specialized_matrices) = matrix.rows.into_iter().fold(
-
+            (vec![], vec![]),
-        let specialized_matrices = row_iter
+            |(mut default_matrix, mut case_matrices): (Vec<Row>, Vec<(CaseTest, Vec<Row>)>),
-            .peeking_take_while(|row| !match_wild_card(&row.columns[occurrence_col].pattern))
+             mut row| {
-            .fold(vec![], |mut case_matrices, mut row| {
+                // For example in the case of matching on []
                 if row.columns.is_empty() {
-                    case_matrices.push((CaseTest::Wild, vec![row]));
+                    default_matrix.push(row);
-                    return case_matrices;
+                    return (default_matrix, case_matrices);
                 }
 
                 let col = row.columns.remove(occurrence_col);
 
                 let (case, remaining_patts) = match col.pattern {
+                    Pattern::Var { .. } | Pattern::Discard { .. } => (CaseTest::Wild, vec![]),
                     Pattern::Int { value, .. } => (CaseTest::Int(value.clone()), vec![]),
                     Pattern::ByteArray { value, .. } => (CaseTest::Bytes(value.clone()), vec![]),
                     Pattern::List { elements, tail, .. } => (
@@ -281,52 +508,73 @@ impl<'a, 'b> TreeGen<'a, 'b> {
                                 .collect_vec(),
                         )
                     }
-                    Pattern::Tuple { .. }
+                    Pattern::Tuple { .. } | Pattern::Pair { .. } | Pattern::Assign { .. } => {
-                    | Pattern::Pair { .. }
-                    | Pattern::Assign { .. }
-                    | Pattern::Var { .. }
-                    | Pattern::Discard { .. } => {
                         unreachable!("{:#?}", col.pattern)
                     }
                 };
 
                 // Assert path is the same for each specialized row
-                assert!(path == col.path);
+                assert!(path == col.path || matches!(case, CaseTest::Wild));
 
                 // expand assigns by newly added ones
                 row.assigns
                     .extend(remaining_patts.iter().flat_map(|x| x.0.clone()));
 
                 // Add inner patterns to existing row
-                row.columns
+                let mut new_cols = remaining_patts.into_iter().flat_map(|x| x.1).collect_vec();
-                    .extend(remaining_patts.into_iter().flat_map(|x| x.1));
 
-                // For lists with tail it's a special case where we also add it to existing patterns
+                let added_columns = new_cols.len();
-                // all the way to the longest element. The reason being that each list size greater
+
-                // than the list with tail could also match with could also match depending on the inner pattern.
+                // Pop off tail so that it aligns more easily with other list patterns
-                // See tests below for an example
+                if matches!(case, CaseTest::ListWithTail(_)) {
-                if let CaseTest::ListWithTail(elems_len) = case {
+                    new_cols.pop();
+                }
+
+                new_cols.extend(row.columns);
+
+                row.columns = new_cols;
+
+                if let CaseTest::Wild = case {
+                    let current_wild_cols = row.columns.len();
+                    default_matrix.push(row.clone());
+
+                    case_matrices.iter_mut().for_each(|(_, matrix)| {
+                        let mut row = row.clone();
+                        let total_cols = matrix[0].columns.len();
+
+                        if total_cols != 0 {
+                            let added_columns = total_cols - current_wild_cols;
+
+                            for _ in 0..added_columns {
+                                row.columns.insert(0, self.wild_card_pattern.clone());
+                            }
+
+                            matrix.push(row);
+                        }
+                    });
+                } else if let CaseTest::ListWithTail(case_length) = case {
+                    // For lists with tail it's a special case where we also add it to existing patterns
+                    // all the way to the longest element. The reason being that each list size greater
+                    // than the list with tail could also match with could also match depending on the inner pattern.
+                    // See tests below for an example
                     if let Some(longest_elems_no_tail) = longest_elems_no_tail {
-                        for elem_count in elems_len..=longest_elems_no_tail {
+                        for elem_count in case_length..=longest_elems_no_tail {
                             let case = CaseTest::List(elem_count);
 
                             let mut row = row.clone();
 
-                            let tail = row.columns.pop().unwrap();
+                            for _ in 0..(elem_count - case_length) {
-
+                                row.columns
-                            let columns_to_fill = (0..(elem_count - elems_len))
+                                    .insert(case_length, self.wild_card_pattern.clone());
-                                .map(|_| tail.clone())
-                                .collect_vec();
-
-                            row.columns.extend(columns_to_fill);
-
-                            if let Some(entry) =
-                                case_matrices.iter_mut().find(|item| item.0 == case)
-                            {
-                                entry.1.push(row);
-                            } else {
-                                case_matrices.push((case, vec![row]));
                             }
+
+                            self.insert_case(
+                                &mut case_matrices,
+                                case,
+                                &default_matrix,
+                                row,
+                                added_columns,
+                            );
                         }
                     }
 
@@ -334,190 +582,138 @@ impl<'a, 'b> TreeGen<'a, 'b> {
                         unreachable!()
                     };
 
-                    for elem_count in elems_len..=longest_elems_with_tail {
+                    for elem_count in case_length..=longest_elems_with_tail {
                         let case = CaseTest::ListWithTail(elem_count);
 
                         let mut row = row.clone();
 
-                        let tail = row.columns.pop().unwrap();
+                        for _ in 0..(elem_count - case_length) {
-
+                            row.columns
-                        let columns_to_fill = (0..(elem_count - elems_len))
+                                .insert(case_length, self.wild_card_pattern.clone());
-                            .map(|_| tail.clone())
-                            .collect_vec();
-
-                        row.columns.extend(columns_to_fill);
-
-                        if let Some(entry) = case_matrices.iter_mut().find(|item| item.0 == case) {
-                            entry.1.push(row);
-                        } else {
-                            case_matrices.push((case, vec![row]));
                         }
+
+                        self.insert_case(
+                            &mut case_matrices,
+                            case,
+                            &default_matrix,
+                            row,
+                            added_columns,
+                        );
                     }
                 } else {
-                    if let Some(entry) = case_matrices.iter_mut().find(|item| item.0 == case) {
+                    self.insert_case(
-                        entry.1.push(row);
+                        &mut case_matrices,
-                    } else {
+                        case,
-                        case_matrices.push((case, vec![row]));
+                        &default_matrix,
-                    }
+                        row,
+                        added_columns,
+                    );
                 }
 
-                case_matrices
+                (default_matrix, case_matrices)
-            });
+            },
+        );
 
         let default_matrix = PatternMatrix {
-            rows: row_iter.collect_vec(),
+            rows: default_matrix,
         };
 
         if has_list_pattern {
             // Since the list_tail case might cover the rest of the possible matches extensively
             // then fallback is optional here
             let fallback_option = if default_matrix.rows.is_empty() {
-                fallback_option
+                None
             } else {
-                Some(self.do_build_tree(
+                Some(
-                    subject_name,
+                    self.do_build_tree(
-                    subject_tipo,
+                        subject_name,
-                    // Since everything after this point had a wild card on or above
+                        subject_tipo,
-                    // the row for the selected column in front. Then we ignore the
+                        // Since everything after this point had a wild card on or above
-                    // cases and continue to check other columns.
+                        // the row for the selected column in front. Then we ignore the
-                    default_matrix,
+                        // cases and continue to check other columns.
-                    fallback_option,
+                        default_matrix,
-                ))
+                        then_map,
+                    )
+                    .into(),
+                )
             };
 
             let (tail_cases, cases): (Vec<_>, Vec<_>) = specialized_matrices
                 .into_iter()
                 .partition(|(case, _)| matches!(case, CaseTest::ListWithTail(_)));
 
-            // TODO: pass in interner and use unique string
+            DecisionTree::ListSwitch {
-            let hoisted_name = "HoistedThing".to_string();
+                subject_name: subject_name.clone(),
-
+                subject_tipo: specialized_tipo.clone(),
-            if let Some(fallback) = fallback_option {
+                path,
-                DecisionTree::HoistThen(
+                cases: cases
-                    hoisted_name.clone(),
+                    .into_iter()
-                    fallback.into(),
+                    .map(|x| {
-                    DecisionTree::ListSwitch {
+                        (
-                        subject_name: subject_name.clone(),
+                            x.0,
-                        subject_tipo: specialized_tipo.clone(),
+                            self.do_build_tree(
-                        path,
+                                subject_name,
-                        cases: cases
+                                subject_tipo,
-                            .into_iter()
+                                PatternMatrix { rows: x.1 },
-                            .map(|x| {
+                                then_map,
-                                (
+                            ),
-                                    x.0,
+                        )
-                                    self.do_build_tree(
+                    })
-                                        subject_name,
+                    .collect_vec(),
-                                        subject_tipo,
+                tail_cases: tail_cases
-                                        PatternMatrix { rows: x.1 },
+                    .into_iter()
-                                        Some(DecisionTree::HoistedLeaf(hoisted_name.clone())),
+                    .map(|x| {
-                                    ),
+                        (
-                                )
+                            x.0,
-                            })
+                            self.do_build_tree(
-                            .collect_vec(),
+                                subject_name,
-                        tail_cases: tail_cases
+                                subject_tipo,
-                            .into_iter()
+                                PatternMatrix { rows: x.1 },
-                            .map(|x| {
+                                then_map,
-                                (
+                            ),
-                                    x.0,
+                        )
-                                    self.do_build_tree(
+                    })
-                                        subject_name,
+                    .collect_vec(),
-                                        subject_tipo,
+                default: fallback_option,
-                                        PatternMatrix { rows: x.1 },
-                                        Some(DecisionTree::HoistedLeaf(hoisted_name.clone())),
-                                    ),
-                                )
-                            })
-                            .collect_vec(),
-                        default: Some(DecisionTree::HoistedLeaf(hoisted_name).into()),
-                    }
-                    .into(),
-                )
-            } else {
-                DecisionTree::ListSwitch {
-                    subject_name: subject_name.clone(),
-                    subject_tipo: specialized_tipo.clone(),
-                    path,
-                    cases: cases
-                        .into_iter()
-                        .map(|x| {
-                            (
-                                x.0,
-                                self.do_build_tree(
-                                    subject_name,
-                                    subject_tipo,
-                                    PatternMatrix { rows: x.1 },
-                                    None,
-                                ),
-                            )
-                        })
-                        .collect_vec(),
-                    tail_cases: tail_cases
-                        .into_iter()
-                        .map(|x| {
-                            (
-                                x.0,
-                                self.do_build_tree(
-                                    subject_name,
-                                    subject_tipo,
-                                    PatternMatrix { rows: x.1 },
-                                    None,
-                                ),
-                            )
-                        })
-                        .collect_vec(),
-                    default: None,
-                }
             }
-        } else if specialized_matrices.is_empty() {
-            // No more patterns to match on so we grab the first default row and return that
-            let mut fallback = default_matrix.rows;
-
-            let row = fallback.swap_remove(0);
-
-            DecisionTree::Leaf(row.assigns, row.then)
         } else {
-            let fallback = if default_matrix.rows.is_empty() {
+            let fallback_option = if default_matrix.rows.is_empty() {
-                fallback_option.unwrap()
+                None
             } else {
-                self.do_build_tree(
+                Some(
-                    subject_name,
+                    self.do_build_tree(
-                    subject_tipo,
+                        subject_name,
-                    // Since everything after this point had a wild card on or above
+                        subject_tipo,
-                    // the row for the selected column in front. Then we ignore the
+                        // Since everything after this point had a wild card on or above
-                    // cases and continue to check other columns.
+                        // the row for the selected column in front. Then we ignore the
-                    default_matrix,
+                        // cases and continue to check other columns.
-                    fallback_option,
+                        default_matrix,
+                        then_map,
+                    )
+                    .into(),
                 )
             };
 
-            // TODO: pass in interner and use unique string
+            DecisionTree::Switch {
-            let hoisted_name = "HoistedThing".to_string();
+                subject_name: subject_name.clone(),
-
+                subject_tipo: specialized_tipo.clone(),
-            DecisionTree::HoistThen(
+                path,
-                hoisted_name.clone(),
+                cases: specialized_matrices
-                fallback.into(),
+                    .into_iter()
-                DecisionTree::Switch {
+                    .map(|x| {
-                    subject_name: subject_name.clone(),
+                        (
-                    subject_tipo: specialized_tipo.clone(),
+                            x.0,
-                    path,
+                            self.do_build_tree(
-                    cases: specialized_matrices
+                                subject_name,
-                        .into_iter()
+                                subject_tipo,
-                        .map(|x| {
+                                PatternMatrix { rows: x.1 },
-                            (
+                                then_map,
-                                x.0,
+                            ),
-                                self.do_build_tree(
+                        )
-                                    subject_name,
+                    })
-                                    subject_tipo,
+                    .collect_vec(),
-                                    PatternMatrix { rows: x.1 },
+                default: fallback_option.into(),
-                                    Some(DecisionTree::HoistedLeaf(hoisted_name.clone())),
+            }
-                                ),
-                            )
-                        })
-                        .collect_vec(),
-                    default: DecisionTree::HoistedLeaf(hoisted_name).into(),
-                }
-                .into(),
-            )
         }
     }
 
@@ -659,6 +855,30 @@ impl<'a, 'b> TreeGen<'a, 'b> {
             }
         }
     }
+
+    fn insert_case(
+        &self,
+        case_matrices: &mut Vec<(CaseTest, Vec<Row<'a>>)>,
+        case: CaseTest,
+        default_matrix: &Vec<Row<'a>>,
+        new_row: Row<'a>,
+        added_columns: usize,
+    ) {
+        if let Some(entry) = case_matrices.iter_mut().find(|item| item.0 == case) {
+            entry.1.push(new_row);
+        } else {
+            let mut rows = default_matrix.clone();
+
+            for _ in 0..added_columns {
+                for row in &mut rows {
+                    row.columns.insert(0, self.wild_card_pattern.clone());
+                }
+            }
+
+            rows.push(new_row);
+            case_matrices.push((case, rows));
+        }
+    }
 }
 
 fn get_tipo_by_path(mut subject_tipo: Rc<Type>, mut path: &[Path]) -> Rc<Type> {
@@ -686,7 +906,8 @@ fn match_wild_card(pattern: &TypedPattern) -> bool {
 }
 
 // A function to get which column has the most pattern matches before a wild card
-fn highest_occurrence(matrix: &PatternMatrix, column_length: usize) -> usize {
+// Returns none if all columns in the first row are wild cards
+fn highest_occurrence(matrix: &PatternMatrix, column_length: usize) -> Option<usize> {
     let occurrences = [Occurrence::default()].repeat(column_length);
 
     let occurrences =
@@ -721,7 +942,11 @@ fn highest_occurrence(matrix: &PatternMatrix, column_length: usize) -> usize {
         }
     });
 
-    highest_occurrence.0
+    if highest_occurrence.1 == 0 {
+        None
+    } else {
+        Some(highest_occurrence.0)
+    }
 }
 
 #[cfg(test)]
@@ -732,7 +957,8 @@ mod tester {
 
     use crate::{
         ast::{
-            well_known, Definition, ModuleKind, TraceLevel, Tracing, TypedModule, UntypedModule,
+            Definition, ModuleKind, Span, TraceLevel, Tracing, TypedModule, TypedPattern,
+            UntypedModule,
         },
         builtins,
         expr::{Type, TypedExpr},
@@ -823,14 +1049,16 @@ mod tester {
 
         let data_types = IndexMap::new();
 
-        let tree_gen = TreeGen {
+        let pattern = TypedPattern::Discard {
-            interner: &mut air_interner,
+            name: "_".to_string(),
-            data_types: &data_types,
+            location: Span::empty(),
         };
 
+        let tree_gen = TreeGen::new(&mut air_interner, &data_types, &pattern);
+
         let tree = tree_gen.build_tree(&"subject".to_string(), &Type::list(Type::int()), clauses);
 
-        println!("TREE IS {:#?}", tree);
+        println!("{:#?}", tree);
     }
 
     #[test]
@@ -859,11 +1087,13 @@ mod tester {
 
         let data_types = IndexMap::new();
 
-        let tree_gen = TreeGen {
+        let pattern = TypedPattern::Discard {
-            interner: &mut air_interner,
+            name: "_".to_string(),
-            data_types: &data_types,
+            location: Span::empty(),
         };
 
+        let tree_gen = TreeGen::new(&mut air_interner, &data_types, &pattern);
+
         let tree = tree_gen.build_tree(
             &"subject".to_string(),
             &Type::tuple(vec![
@@ -875,7 +1105,7 @@ mod tester {
             clauses,
         );
 
-        println!("TREE IS {:#?}", tree);
+        println!("{:#?}", tree);
     }
 
     #[test]
@@ -906,11 +1136,13 @@ mod tester {
 
         let data_types = IndexMap::new();
 
-        let tree_gen = TreeGen {
+        let pattern = TypedPattern::Discard {
-            interner: &mut air_interner,
+            name: "_".to_string(),
-            data_types: &data_types,
+            location: Span::empty(),
         };
 
+        let tree_gen = TreeGen::new(&mut air_interner, &data_types, &pattern);
+
         let tree = tree_gen.build_tree(
             &"subject".to_string(),
             &Type::tuple(vec![
@@ -922,7 +1154,7 @@ mod tester {
             clauses,
         );
 
-        println!("TREE IS {:#?}", tree);
+        println!("{:#?}", tree);
     }
 
     #[test]
@@ -954,11 +1186,13 @@ mod tester {
 
         let data_types = IndexMap::new();
 
-        let tree_gen = TreeGen {
+        let pattern = TypedPattern::Discard {
-            interner: &mut air_interner,
+            name: "_".to_string(),
-            data_types: &data_types,
+            location: Span::empty(),
         };
 
+        let tree_gen = TreeGen::new(&mut air_interner, &data_types, &pattern);
+
         let tree = tree_gen.build_tree(
             &"subject".to_string(),
             &Type::tuple(vec![
@@ -970,7 +1204,8 @@ mod tester {
             clauses,
         );
 
-        println!("TREE IS {:#?}", tree);
+        println!("{}", tree);
+        panic!("SUPPPPPPPPPPPPPPPPPPPPPPPER DOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOONE");
     }
 
     #[test]
@@ -994,7 +1229,7 @@ mod tester {
             panic!()
         };
 
-        let TypedExpr::When { clauses, .. } = &function.body else {
+        let TypedExpr::When { clauses, tipo, .. } = &function.body else {
             panic!()
         };
 
@@ -1004,11 +1239,15 @@ mod tester {
 
         let data_types = builtins::prelude_data_types(&id_gen);
 
-        let tree_gen = TreeGen {
+        let pattern = TypedPattern::Discard {
-            interner: &mut air_interner,
+            name: "_".to_string(),
-            data_types: &utils::indexmap::as_ref_values(&data_types),
+            location: Span::empty(),
         };
 
+        let data_types = utils::indexmap::as_ref_values(&data_types);
+
+        let tree_gen = TreeGen::new(&mut air_interner, &data_types, &pattern);
+
         let tree = tree_gen.build_tree(
             &"subject".to_string(),
             &Type::tuple(vec![
@@ -1020,7 +1259,7 @@ mod tester {
             clauses,
         );
 
-        println!("TREE IS {:#?}", tree);
+        println!("{}", tree);
         panic!("SUPPPPPPPPPPPPPPPPPPPPPPPER DOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOONE");
     }
 }