Add case constr for applies greater than 2 optimization

crates/aiken-project/src/snapshots/aiken_project__export__tests__recursive_types.snap

@@ -24,8 +24,8 @@ description: "Code:\n\npub type Foo<a> {\n  Empty\n  Bar(a, Foo<a>)\n}\n\npub fn
       "$ref": "#/definitions/Int"
     }
   },
-  "compiledCode": "59017d0101003333332222222232323232325333008300430093754002264a666012600a60146ea800452000132337006eb4c038004cc011300103d8798000300e300f001300b37540026018601a00a264a66601266e1d2002300a37540022646466e00cdc01bad300f002375a601e0026600a601e6020004601e602000260186ea8008c02cdd500109919b80375a601c00266008601c601e002980103d8798000300b37540046018601a00a601600860020024446464a666014600c60166ea80044c94ccc02cc01cc030dd50008a400026466e00dd69808000999803803a60103d879800030103011001300d3754002601c601e004264a66601666e1d2002300c37540022646466e00cdc01bad3011002375a60220026660100106022602400460226024002601c6ea8008c034dd500109919b80375a602000266600e00e60206022002980103d8798000300d3754004601c601e004601a002660160046601600297ae0370e90001980300119803000a5eb815d12ba15740aae7955ceab9a01",
+  "compiledCode": "59017c0101009800aba2aba1aba0aab9eaab9dab9a488888888c8c8c8c8c94ccc020c010c024dd50008992999804980298051baa0011480004c8cdc01bad300e001330044c0103d8798000300e300f001300b37540026018601a00a264a66601266e1d2002300a37540022646466e00cdc01bad300f002375a601e0026600a601e6020004601e602000260186ea8008c02cdd500109919b80375a601c00266008601c601e002980103d8798000300b37540046018601a00a601600860020024446464a666014600c60166ea80044c94ccc02cc01cc030dd50008a400026466e00dd69808000999803803a60103d879800030103011001300d3754002601c601e004264a66601666e1d2002300c37540022646466e00cdc01bad3011002375a60220026660100106022602400460226024002601c6ea8008c034dd500109919b80375a602000266600e00e60206022002980103d8798000300d3754004601c601e004601a002660160046601600297ae0370e90001980300119803000a5eb81",
-  "hash": "1f7ba1be8ac6bba7b61d818e0f274b1feae61d24737dd9e833d59430",
+  "hash": "cf7584d2ab87ed23b3443146bbef954c61f15eaf1511295703fe1ccd",
   "definitions": {
     "Int": {
       "dataType": "integer"

crates/uplc/src/builder.rs

@@ -30,6 +30,17 @@ where
         Term::Delay(self.into())
     }
 
+    pub fn constr(tag: usize, fields: Vec<Term<T>>) -> Self {
+        Term::Constr { tag, fields }
+    }
+
+    pub fn case(self, branches: Vec<Term<T>>) -> Self {
+        Term::Case {
+            constr: self.into(),
+            branches,
+        }
+    }
+
     // Primitives
     pub fn integer(i: num_bigint::BigInt) -> Self {
         Term::Constant(Constant::Integer(i).into())

crates/uplc/src/optimize/shrinker.rs

@@ -1362,6 +1362,62 @@ impl Term<Name> {
         }
     }
 
+    // IMPORTANT: RUNS ONE TIME AND ONLY ON THE LAST PASS
+    fn case_constr_apply_reducer(
+        &mut self,
+        _id: Option<usize>,
+        _arg_stack: Vec<Args>,
+        _scope: &Scope,
+        _context: &mut Context,
+    ) {
+        let mut term = &mut std::mem::replace(self, Term::Error.force());
+
+        let mut arg_vec = vec![];
+
+        while let Term::Apply { function, argument } = term {
+            arg_vec.push(Rc::make_mut(argument));
+
+            term = Rc::make_mut(function);
+        }
+
+        arg_vec.reverse();
+
+        match term {
+            Term::Case { constr, branches }
+                if branches.len() == 1 && matches!(constr.as_ref(), Term::Constr { .. }) =>
+            {
+                let Term::Constr { fields, .. } = Rc::make_mut(constr) else {
+                    unreachable!();
+                };
+
+                for arg in arg_vec {
+                    fields.push(std::mem::replace(arg, Term::Error.force()));
+                }
+
+                *self = std::mem::replace(term, Term::Error.force());
+            }
+            _ => {
+                if arg_vec.len() > 3 {
+                    let mut fields = vec![];
+
+                    for arg in arg_vec {
+                        fields.push(std::mem::replace(arg, Term::Error.force()));
+                    }
+
+                    *self = Term::constr(0, fields)
+                        .case(vec![std::mem::replace(term, Term::Error.force())]);
+                } else {
+                    for arg in arg_vec {
+                        *term = (std::mem::replace(term, Term::Error.force()))
+                            .apply(std::mem::replace(arg, Term::Error.force()));
+                    }
+
+                    *self = std::mem::replace(term, Term::Error.force());
+                }
+            }
+        }
+    }
+
     fn identity_reducer(
         &mut self,
         _id: Option<usize>,
@@ -1754,7 +1810,7 @@ impl Term<Name> {
                                 acc.apply(arg.pierce_no_inlines().clone())
                             });
 
-                    // Check above for is error safe
+                    // The check above is to make sure the program is error safe
                     let eval_term: Term<Name> = Program {
                         version: (1, 0, 0),
                         term: applied_term,
@@ -1881,7 +1937,7 @@ impl Program<Name> {
             context,
         )
     }
-    // This one runs the optimizations that are only done a single time
+    // This runs the optimizations that are only done a single time
     pub fn run_once_pass(self) -> Self {
         let program = self
             .traverse_uplc_with(false, &mut |id, term, _arg_stack, scope, context| {
@@ -1946,38 +2002,36 @@ impl Program<Name> {
 
     pub fn multi_pass(self) -> (Self, Context) {
         self.traverse_uplc_with(true, &mut |id, term, arg_stack, scope, context| {
-            let mut changed;
+            let false = term.lambda_reducer(id, arg_stack.clone(), scope, context) else {
+                term.remove_inlined_ids(id, vec![], scope, context);
+                return;
+            };
 
-            changed = term.lambda_reducer(id, arg_stack.clone(), scope, context);
+            let false = term.identity_reducer(id, arg_stack.clone(), scope, context) else {
-            if changed {
                 term.remove_inlined_ids(id, vec![], scope, context);
                 return;
-            }
+            };
-            changed = term.identity_reducer(id, arg_stack.clone(), scope, context);
+
-            if changed {
+            let false = term.inline_reducer(id, arg_stack.clone(), scope, context) else {
                 term.remove_inlined_ids(id, vec![], scope, context);
                 return;
-            }
+            };
-            changed = term.inline_reducer(id, arg_stack.clone(), scope, context);
+
-            if changed {
+            let false = term.force_delay_reducer(id, arg_stack.clone(), scope, context) else {
                 term.remove_inlined_ids(id, vec![], scope, context);
                 return;
-            }
+            };
-            changed = term.force_delay_reducer(id, arg_stack.clone(), scope, context);
+
-            if changed {
+            let false = term.cast_data_reducer(id, arg_stack.clone(), scope, context) else {
                 term.remove_inlined_ids(id, vec![], scope, context);
                 return;
-            }
+            };
-            changed = term.cast_data_reducer(id, arg_stack.clone(), scope, context);
+
-            if changed {
+            let false = term.builtin_eval_reducer(id, arg_stack.clone(), scope, context) else {
                 term.remove_inlined_ids(id, vec![], scope, context);
                 return;
-            }
+            };
-            changed = term.builtin_eval_reducer(id, arg_stack.clone(), scope, context);
+
-            if changed {
-                term.remove_inlined_ids(id, vec![], scope, context);
-                return;
-            }
             term.convert_arithmetic_ops(id, arg_stack, scope, context);
             term.flip_constants(id, vec![], scope, context);
             term.remove_inlined_ids(id, vec![], scope, context);
@@ -2000,6 +2054,7 @@ impl Program<Name> {
     pub fn clean_up(self) -> Self {
         self.traverse_uplc_with(true, &mut |id, term, _arg_stack, scope, context| {
             term.remove_no_inlines(id, vec![], scope, context);
+            term.case_constr_apply_reducer(id, vec![], scope, context);
         })
         .0
     }
@@ -2043,7 +2098,6 @@ impl Program<Name> {
                             ) {
                             // We found it the builtin was curried before
                             // So now we merge the new args into the existing curried builtin
-
                             let curried_builtin = curried_terms.swap_remove(index);
 
                             let curried_builtin =
@@ -3403,4 +3457,63 @@ mod tests {
 
         compare_optimization(expected, program, |p| p.builtin_curry_reducer());
     }
+
+    #[test]
+    fn case_constr_apply_test_1() {
+        let program: Program<Name> = Program {
+            version: (1, 1, 0),
+            term: Term::add_integer()
+                .apply(Term::integer(0.into()))
+                .apply(Term::integer(0.into()))
+                .apply(Term::integer(0.into()))
+                .apply(Term::integer(0.into()))
+                .apply(Term::integer(0.into()))
+                .apply(Term::integer(0.into())),
+        };
+
+        let expected = Program {
+            version: (1, 1, 0),
+            term: Term::constr(
+                0,
+                vec![
+                    Term::integer(0.into()),
+                    Term::integer(0.into()),
+                    Term::integer(0.into()),
+                    Term::integer(0.into()),
+                    Term::integer(0.into()),
+                    Term::integer(0.into()),
+                ],
+            )
+            .case(vec![Term::add_integer()]),
+        };
+
+        compare_optimization(expected, program, |p| {
+            p.run_one_opt(true, &mut |id, term, arg_stack, scope, context| {
+                term.case_constr_apply_reducer(id, arg_stack, scope, context);
+            })
+        });
+    }
+
+    #[test]
+    fn case_constr_apply_test_2() {
+        let program: Program<Name> = Program {
+            version: (1, 1, 0),
+            term: Term::add_integer()
+                .apply(Term::integer(0.into()))
+                .apply(Term::integer(0.into())),
+        };
+
+        let expected = Program {
+            version: (1, 1, 0),
+            term: Term::add_integer()
+                .apply(Term::integer(0.into()))
+                .apply(Term::integer(0.into())),
+        };
+
+        compare_optimization(expected, program, |p| {
+            p.run_one_opt(true, &mut |id, term, arg_stack, scope, context| {
+                term.case_constr_apply_reducer(id, arg_stack, scope, context);
+            })
+        });
+    }
 }