Inline now handles (if cond then body else error) patterns.

This allows conditions like ```expect x == 1``` to match performance with ```x == 1 && ...```

Also change builtins forcing to accommodate the new case-constr apply optimization

crates/aiken-lang/src/gen_uplc.rs

@@ -217,9 +217,11 @@ impl<'a> CodeGenerator<'a> {
 
     fn finalize(&mut self, mut term: Term<Name>) -> Program<Name> {
         term = self.special_functions.apply_used_functions(term);
-
+        println!("PROG BEFORE IS  {}", term.to_pretty());
         let program = aiken_optimize_and_intern(self.new_program(term));
 
+        println!("PROG IS {}", program.to_pretty());
+
         // This is very important to call here.
         // If this isn't done, re-using the same instance
         // of the generator will result in free unique errors

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

@@ -1,7 +1,6 @@
 ---
 source: crates/aiken-project/src/export.rs
 description: "Code:\n\npub type Foo<a> {\n  Empty\n  Bar(a, Foo<a>)\n}\n\npub fn add(a: Foo<Int>, b: Foo<Int>) -> Int {\n  when (a, b) is {\n    (Empty, Empty) -> 0\n    (Bar(x, y), Bar(c, d)) -> x + c + add(y, d)\n    (Empty, Bar(c, d)) -> c + add(Empty, d)\n    (Bar(x, y), Empty) -> x + add(y, Empty)\n  }\n}\n"
-snapshot_kind: text
 ---
 {
   "name": "test_module.add",
@@ -25,8 +24,8 @@ snapshot_kind: text
       "$ref": "#/definitions/Int"
     }
   },
-  "compiledCode": "59017d0101003232323232322232323232325333008300430093754002264a666012600a60146ea800452000132337006eb4c038004cc011300103d8798000300e300f001300b37540026018601a00a264a66601266e1d2002300a37540022646466e00cdc01bad300f002375a601e0026600a601e6020004601e602000260186ea8008c02cdd500109919b80375a601c00266008601c601e002980103d8798000300b37540046018601a00a601600860020024446464a666014600c60166ea80044c94ccc02cc01cc030dd50008a400026466e00dd69808000999803803a60103d879800030103011001300d3754002601c601e004264a66601666e1d2002300c37540022646466e00cdc01bad3011002375a60220026660100106022602400460226024002601c6ea8008c034dd500109919b80375a602000266600e00e60206022002980103d8798000300d3754004601c601e004601a002660160046601600297ae0370e90001980300119803000a5eb815cd2ab9d5573cae815d0aba201",
+  "compiledCode": "59017d0101003333332222222232323232325333008300430093754002264a666012600a60146ea800452000132337006eb4c038004cc011300103d8798000300e300f001300b37540026018601a00a264a66601266e1d2002300a37540022646466e00cdc01bad300f002375a601e0026600a601e6020004601e602000260186ea8008c02cdd500109919b80375a601c00266008601c601e002980103d8798000300b37540046018601a00a601600860020024446464a666014600c60166ea80044c94ccc02cc01cc030dd50008a400026466e00dd69808000999803803a60103d879800030103011001300d3754002601c601e004264a66601666e1d2002300c37540022646466e00cdc01bad3011002375a60220026660100106022602400460226024002601c6ea8008c034dd500109919b80375a602000266600e00e60206022002980103d8798000300d3754004601c601e004601a002660160046601600297ae0370e90001980300119803000a5eb815d12ba15740aae7955ceab9a01",
-  "hash": "c6af3f04e300cb8c1d0429cc0d8e56a0413eef9fcb338f72076b426c",
+  "hash": "1f7ba1be8ac6bba7b61d818e0f274b1feae61d24737dd9e833d59430",
   "definitions": {
     "Int": {
       "dataType": "integer"

crates/aiken-project/src/tests/gen_uplc.rs

@@ -64,6 +64,7 @@ fn assert_uplc(source_code: &str, expected: Term<Name>, should_fail: bool, verbo
                 version: (1, 1, 0),
                 term: expected,
             };
+            println!("BEFORE OPT IS {}", expected.to_pretty());
 
             let expected = optimize::aiken_optimize_and_intern(expected);
 
@@ -3603,7 +3604,7 @@ fn when_bool_is_true() {
     assert_uplc(
         src,
         Term::var("subject")
-            .delayed_if_then_else(Term::bool(true), Term::Error)
+            .delayed_if_then_else(Term::bool(true), Term::Error.delay().force())
             .lambda("subject")
             .apply(Term::bool(true)),
         false,
@@ -3627,7 +3628,7 @@ fn when_bool_is_true_switched_cases() {
     assert_uplc(
         src,
         Term::var("subject")
-            .delayed_if_then_else(Term::bool(true), Term::Error)
+            .delayed_if_then_else(Term::bool(true), Term::Error.delay().force())
             .lambda("subject")
             .apply(Term::bool(true)),
         false,
@@ -3651,7 +3652,7 @@ fn when_bool_is_false() {
     assert_uplc(
         src,
         Term::var("subject")
-            .delayed_if_then_else(Term::bool(true), Term::Error)
+            .delayed_if_then_else(Term::bool(true), Term::Error.delay().force())
             .lambda("subject")
             .apply(Term::bool(false)),
         true,
@@ -4088,17 +4089,17 @@ fn generic_validator_type_test() {
                         Term::tail_list()
                             .apply(Term::Var(tail_id_5.clone()))
                             .as_var("tail_id_6", |tail_id_6| {
-                                Term::head_list()
-                                    .apply(Term::Var(tail_id_6.clone()))
-                                    .as_var("__val", |val| {
                                 Term::tail_list()
-                                            .apply(Term::Var(tail_id_6))
+                                    .apply(Term::Var(tail_id_6.clone()))
                                     .delayed_choose_list(
-                                                expect_b(val, Term::Var(then_delayed), trace),
+                                        Term::head_list()
+                                            .apply(Term::Var(tail_id_6))
+                                            .as_var("__val", |val| {
+                                                expect_b(val, Term::Var(then_delayed), trace)
+                                            }),
                                         Term::Error,
                                     )
                             })
-                            })
                     }
                 });
 

crates/uplc/src/optimize/shrinker.rs

@@ -167,6 +167,7 @@ impl DefaultFunction {
         )
     }
     /// For now all of the curry builtins are not forceable
+    /// Curryable builtins must take in 2 or more arguments
     pub fn can_curry_builtin(self) -> bool {
         matches!(
             self,
@@ -205,7 +206,8 @@ impl DefaultFunction {
             | DefaultFunction::MultiplyInteger
             | DefaultFunction::EqualsInteger
             | DefaultFunction::LessThanInteger
-            | DefaultFunction::LessThanEqualsInteger => arg_stack.iter().all(|arg| {
+            | DefaultFunction::LessThanEqualsInteger
+            | DefaultFunction::IData => arg_stack.iter().all(|arg| {
                 if let Term::Constant(c) = arg {
                     matches!(c.as_ref(), Constant::Integer(_))
                 } else {
@@ -226,10 +228,13 @@ impl DefaultFunction {
                     false
                 }
             }),
-            DefaultFunction::EqualsByteString
+            DefaultFunction::LengthOfByteString
+            | DefaultFunction::EqualsByteString
             | DefaultFunction::AppendByteString
             | DefaultFunction::LessThanEqualsByteString
-            | DefaultFunction::LessThanByteString => arg_stack.iter().all(|arg| {
+            | DefaultFunction::LessThanByteString
+            | DefaultFunction::DecodeUtf8
+            | DefaultFunction::BData => arg_stack.iter().all(|arg| {
                 if let Term::Constant(c) = arg {
                     matches!(c.as_ref(), Constant::ByteString(_))
                 } else {
@@ -282,23 +287,25 @@ impl DefaultFunction {
                 }
             }
 
-            DefaultFunction::EqualsString | DefaultFunction::AppendString => {
+            DefaultFunction::EqualsString
-                arg_stack.iter().all(|arg| {
+            | DefaultFunction::AppendString
+            | DefaultFunction::EncodeUtf8 => arg_stack.iter().all(|arg| {
                 if let Term::Constant(c) = arg {
                     matches!(c.as_ref(), Constant::String(_))
                 } else {
                     false
                 }
-                })
+            }),
-            }
 
-            DefaultFunction::EqualsData => arg_stack.iter().all(|arg| {
+            DefaultFunction::EqualsData | DefaultFunction::SerialiseData => {
+                arg_stack.iter().all(|arg| {
                     if let Term::Constant(c) = arg {
                         matches!(c.as_ref(), Constant::Data(_))
                     } else {
                         false
                     }
-            }),
+                })
+            }
 
             DefaultFunction::Bls12_381_G1_Equal | DefaultFunction::Bls12_381_G1_Add => {
                 arg_stack.iter().all(|arg| {
@@ -337,6 +344,10 @@ impl DefaultFunction {
             _ => false,
         }
     }
+
+    pub fn wrapped_name(self) -> String {
+        format!("__{}_wrapped", self.aiken_name())
+    }
 }
 
 #[derive(PartialEq, Clone, Debug)]
@@ -1102,7 +1113,7 @@ impl Term<Name> {
                 }
             }
             Term::Delay(body) => {
-                let not_forced: isize = isize::from(force_stack.pop().is_none());
+                let not_forced = isize::from(force_stack.pop().is_none());
 
                 body.var_occurrences(search_for, arg_stack, force_stack)
                     .delay_if_found(not_forced)
@@ -1125,11 +1136,34 @@ impl Term<Name> {
                 }
             }
             Term::Apply { function, argument } => {
+                // unwrap apply and add void to arg stack!
                 arg_stack.push(());
 
-                function
+                let apply_var_occurrence_stack = |term: &Term<Name>, arg_stack: Vec<()>| {
-                    .var_occurrences(search_for.clone(), arg_stack, force_stack)
+                    term.var_occurrences(search_for.clone(), arg_stack, force_stack)
-                    .combine(argument.var_occurrences(search_for, vec![], vec![]))
+                };
+
+                let apply_var_occurrence_no_stack =
+                    |term: &Term<Name>| term.var_occurrences(search_for.clone(), vec![], vec![]);
+
+                if let Term::Apply {
+                    function: next_func,
+                    argument: next_arg,
+                } = function.as_ref()
+                {
+                    // unwrap apply and add void to arg stack!
+                    arg_stack.push(());
+                    next_func.carry_args_to_branch(
+                        next_arg,
+                        argument,
+                        arg_stack,
+                        apply_var_occurrence_stack,
+                        apply_var_occurrence_no_stack,
+                    )
+                } else {
+                    apply_var_occurrence_stack(function, arg_stack)
+                        .combine(apply_var_occurrence_no_stack(argument))
+                }
             }
             Term::Force(x) => {
                 force_stack.push(());
@@ -1141,6 +1175,81 @@ impl Term<Name> {
         }
     }
 
+    // This handles the very common case of (if condition then body else error)
+    // or (if condition then error else body)
+    // In this case it is fine to treat the body as if it is not delayed
+    // since the other branch is error
+    fn carry_args_to_branch(
+        &self,
+        then_arg: &Rc<Term<Name>>,
+        else_arg: &Rc<Term<Name>>,
+        mut arg_stack: Vec<()>,
+        var_occurrence_stack: impl FnOnce(&Term<Name>, Vec<()>) -> VarLookup,
+        var_occurrence_no_stack: impl Fn(&Term<Name>) -> VarLookup,
+    ) -> VarLookup {
+        let Term::Apply {
+            function: builtin,
+            argument: condition,
+        } = self
+        else {
+            return var_occurrence_stack(self, arg_stack)
+                .combine(var_occurrence_no_stack(then_arg))
+                .combine(var_occurrence_no_stack(else_arg));
+        };
+
+        // unwrap apply and add void to arg stack!
+        arg_stack.push(());
+
+        let Term::Delay(else_arg) = else_arg.as_ref() else {
+            return var_occurrence_stack(builtin, arg_stack)
+                .combine(var_occurrence_no_stack(condition))
+                .combine(var_occurrence_no_stack(then_arg))
+                .combine(var_occurrence_no_stack(else_arg));
+        };
+
+        let Term::Delay(then_arg) = then_arg.as_ref() else {
+            return var_occurrence_stack(builtin, arg_stack)
+                .combine(var_occurrence_no_stack(condition))
+                .combine(var_occurrence_no_stack(then_arg))
+                .combine(var_occurrence_no_stack(else_arg));
+        };
+
+        match builtin.as_ref() {
+            Term::Var(a)
+                if a.text == DefaultFunction::IfThenElse.wrapped_name()
+                    || a.text == DefaultFunction::ChooseList.wrapped_name() =>
+            {
+                if matches!(else_arg.as_ref(), Term::Error) {
+                    // Pop 3 args of arg_stack due to branch execution
+                    arg_stack.pop();
+                    arg_stack.pop();
+                    arg_stack.pop();
+
+                    var_occurrence_no_stack(condition)
+                        .combine(var_occurrence_stack(then_arg, arg_stack))
+                } else if matches!(then_arg.as_ref(), Term::Error) {
+                    // Pop 3 args of arg_stack due to branch execution
+                    arg_stack.pop();
+                    arg_stack.pop();
+                    arg_stack.pop();
+
+                    var_occurrence_no_stack(condition)
+                        .combine(var_occurrence_stack(else_arg, arg_stack))
+                } else {
+                    var_occurrence_stack(builtin, arg_stack)
+                        .combine(var_occurrence_no_stack(condition))
+                        .combine(var_occurrence_no_stack(then_arg))
+                        .combine(var_occurrence_no_stack(else_arg))
+                }
+            }
+
+            _ => var_occurrence_stack(builtin, arg_stack)
+                .combine(var_occurrence_no_stack(condition))
+                .combine(var_occurrence_no_stack(then_arg))
+                .combine(var_occurrence_no_stack(else_arg)),
+        }
+    }
+
     fn lambda_reducer(
         &mut self,
         _id: Option<usize>,
@@ -1207,7 +1316,7 @@ impl Term<Name> {
 
             if has_forces {
                 context.builtins_map.insert(*func as u8, ());
-                *self = Term::var(format!("__{}_wrapped", func.aiken_name()));
+                *self = Term::var(func.wrapped_name());
             }
         }
     }
@@ -1269,25 +1378,19 @@ impl Term<Name> {
             } => {
                 let body = Rc::make_mut(body);
                 // pops stack here no matter what
-                let temp = Term::Error;
 
-                if let (
+                let Some(Args::Apply(arg_id, identity_func)) = arg_stack.pop() else {
-                    arg_id,
+                    return false;
-                    Term::Lambda {
+                };
+
+                let Term::Lambda {
                     parameter_name: identity_name,
                     body: identity_body,
-                    },
+                } = identity_func.pierce_no_inlines()
-                ) = match &arg_stack.pop() {
+                else {
-                    Some(Args::Apply(
+                    return false;
-                        arg_id,
+                };
-                        Term::Lambda {
+
-                            parameter_name: inline_name,
-                            body,
-                        },
-                    )) if inline_name.text == NO_INLINE => (*arg_id, body.as_ref()),
-                    Some(Args::Apply(arg_id, term)) => (*arg_id, term),
-                    _ => (0, &temp),
-                } {
                 let Term::Var(identity_var) = identity_body.as_ref() else {
                     return false;
                 };
@@ -1309,7 +1412,6 @@ impl Term<Name> {
                     }
                 }
             }
-            }
             Term::Constr { .. } => todo!(),
             Term::Case { .. } => todo!(),
             _ => (),
@@ -1333,22 +1435,20 @@ impl Term<Name> {
                 body,
             } => {
                 // pops stack here no matter what
-                if let Some(Args::Apply(arg_id, arg_term)) = arg_stack.pop() {
+                let Some(Args::Apply(arg_id, arg_term)) = arg_stack.pop() else {
-                    let arg_term = match &arg_term {
+                    return false;
-                        Term::Lambda {
-                            parameter_name,
-                            body,
-                        } if parameter_name.text == NO_INLINE => body.as_ref().clone(),
-                        _ => arg_term,
                 };
 
+                let arg_term = arg_term.pierce_no_inlines();
+
                 let body = Rc::make_mut(body);
 
                 let var_lookup = body.var_occurrences(parameter_name.clone(), vec![], vec![]);
 
-                    let substitute_condition = (var_lookup.delays == 0 && !var_lookup.no_inline)
+                let must_execute_condition = var_lookup.delays == 0 && !var_lookup.no_inline;
-                        || matches!(
+
-                            &arg_term,
+                let cant_throw_condition = matches!(
+                    arg_term,
                     Term::Var(_)
                         | Term::Constant(_)
                         | Term::Delay(_)
@@ -1356,30 +1456,23 @@ impl Term<Name> {
                         | Term::Builtin(_),
                 );
 
-                    if var_lookup.occurrences == 1 && substitute_condition {
+                // This will inline terms that only occur once
+                // if they are guaranteed to execute or can't throw an error by themselves
+                if var_lookup.occurrences == 1 && (must_execute_condition || cant_throw_condition) {
                     changed = true;
-                        body.substitute_var(parameter_name.clone(), arg_term.pierce_no_inlines());
+                    body.substitute_var(parameter_name.clone(), arg_term);
 
                     context.inlined_apply_ids.push(arg_id);
                     *self = std::mem::replace(body, Term::Error.force());
 
                 // This will strip out unused terms that can't throw an error by themselves
-                    } else if !var_lookup.found
+                } else if !var_lookup.found && cant_throw_condition {
-                        && matches!(
-                            arg_term,
-                            Term::Var(_)
-                                | Term::Constant(_)
-                                | Term::Delay(_)
-                                | Term::Lambda { .. }
-                                | Term::Builtin(_)
-                        )
-                    {
                     changed = true;
                     context.inlined_apply_ids.push(arg_id);
                     *self = std::mem::replace(body, Term::Error.force());
                 }
             }
-            }
+
             Term::Constr { .. } => todo!(),
             Term::Case { .. } => todo!(),
             _ => {}
@@ -1647,10 +1740,7 @@ impl Term<Name> {
                         term.pierce_no_inlines()
                     })
                     .collect_vec();
-                if func.can_curry_builtin()
+                if arg_stack.len() == func.arity() && func.is_error_safe(&args) {
-                    && arg_stack.len() == func.arity()
-                    && func.is_error_safe(&args)
-                {
                     changed = true;
                     let applied_term =
                         arg_stack
@@ -1671,7 +1761,7 @@ impl Term<Name> {
                     }
                     .to_named_debruijn()
                     .unwrap()
-                    .eval(ExBudget::max())
+                    .eval(ExBudget::default())
                     .result()
                     .unwrap()
                     .try_into()
@@ -1827,9 +1917,13 @@ impl Program<Name> {
         for default_func_index in context.builtins_map.keys().sorted().cloned() {
             let default_func: DefaultFunction = default_func_index.try_into().unwrap();
 
-            term = term
+            term = term.lambda(default_func.wrapped_name());
-                .lambda(format!("__{}_wrapped", default_func.aiken_name()))
+        }
-                .apply(if default_func.force_count() == 1 {
+
+        for default_func_index in context.builtins_map.keys().sorted().cloned().rev() {
+            let default_func: DefaultFunction = default_func_index.try_into().unwrap();
+
+            term = term.apply(if default_func.force_count() == 1 {
                 Term::Builtin(default_func).force()
             } else {
                 Term::Builtin(default_func).force().force()
@@ -2165,10 +2259,11 @@ fn is_a_builtin_wrapper(term: &Term<Name>) -> bool {
 
     while let Term::Apply { function, argument } = term {
         match argument.as_ref() {
-            Term::Var(name) => arg_names.push(name),
+            Term::Var(name) => arg_names.push(format!("{}_{}", name.text, name.unique)),
 
             Term::Constant(_) => {}
             _ => {
+                //Break loop, it's not a builtin wrapper function
                 return false;
             }
         }
@@ -2178,7 +2273,7 @@ fn is_a_builtin_wrapper(term: &Term<Name>) -> bool {
     arg_names.iter().all(|item| names.contains(item)) && matches!(term, Term::Builtin(_))
 }
 
-fn pop_lambdas_and_get_names(term: &Term<Name>) -> (Vec<Rc<Name>>, &Term<Name>) {
+fn pop_lambdas_and_get_names(term: &Term<Name>) -> (Vec<String>, &Term<Name>) {
     let mut names = vec![];
 
     let mut term = term;
@@ -2189,7 +2284,7 @@ fn pop_lambdas_and_get_names(term: &Term<Name>) -> (Vec<Rc<Name>>, &Term<Name>)
     } = term
     {
         if parameter_name.text != NO_INLINE {
-            names.push(parameter_name.clone());
+            names.push(format!("{}_{}", parameter_name.text, parameter_name.unique));
         }
         term = body.as_ref();
     }
@@ -2365,11 +2460,11 @@ mod tests {
                 .lambda("y")
                 // Forces are automatically applied by builder
                 .lambda("__cons_list_wrapped")
-                .apply(Term::mk_cons())
                 .lambda("__head_list_wrapped")
-                .apply(Term::head_list())
                 .lambda("__tail_list_wrapped")
-                .apply(Term::tail_list()),
+                .apply(Term::tail_list())
+                .apply(Term::head_list())
+                .apply(Term::mk_cons()),
         };
 
         compare_optimization(expected, program, |p| p.run_once_pass());
@@ -2453,9 +2548,9 @@ mod tests {
                         .apply(Term::data(Data::integer(5.into()))),
                 )
                 .lambda("__fst_pair_wrapped")
-                .apply(Term::fst_pair())
                 .lambda("__snd_pair_wrapped")
-                .apply(Term::snd_pair()),
+                .apply(Term::snd_pair())
+                .apply(Term::fst_pair()),
         };
 
         compare_optimization(expected, program, |p| p.run_once_pass());
@@ -2644,6 +2739,80 @@ mod tests {
         });
     }
 
+    #[test]
+    fn inline_reduce_if_then_else_then() {
+        let program: Program<Name> = Program {
+            version: (1, 0, 0),
+            term: Term::var("__if_then_else_wrapped")
+                .apply(Term::bool(true))
+                .apply(Term::sha3_256().apply(Term::var("x")).delay())
+                .apply(Term::Error.delay())
+                .force()
+                .lambda("x")
+                .apply(Term::sha3_256().apply(Term::byte_string(vec![])))
+                .lambda("__if_then_else_wrapped")
+                .apply(Term::Builtin(DefaultFunction::IfThenElse).force()),
+        };
+
+        let expected = Program {
+            version: (1, 0, 0),
+            term: Term::var("__if_then_else_wrapped")
+                .apply(Term::bool(true))
+                .apply(
+                    Term::sha3_256()
+                        .apply(Term::sha3_256().apply(Term::byte_string(vec![])))
+                        .delay(),
+                )
+                .apply(Term::Error.delay())
+                .force()
+                .lambda("__if_then_else_wrapped")
+                .apply(Term::Builtin(DefaultFunction::IfThenElse).force()),
+        };
+
+        compare_optimization(expected, program, |p| {
+            p.run_one_opt(true, &mut |id, term, arg_stack, scope, context| {
+                term.inline_reducer(id, arg_stack, scope, context);
+            })
+        });
+    }
+
+    #[test]
+    fn inline_reduce_if_then_else_else() {
+        let program: Program<Name> = Program {
+            version: (1, 0, 0),
+            term: Term::var("__if_then_else_wrapped")
+                .apply(Term::bool(true))
+                .apply(Term::Error.delay())
+                .apply(Term::sha3_256().apply(Term::var("x")).delay())
+                .force()
+                .lambda("x")
+                .apply(Term::sha3_256().apply(Term::byte_string(vec![])))
+                .lambda("__if_then_else_wrapped")
+                .apply(Term::Builtin(DefaultFunction::IfThenElse).force()),
+        };
+
+        let expected = Program {
+            version: (1, 0, 0),
+            term: Term::var("__if_then_else_wrapped")
+                .apply(Term::bool(true))
+                .apply(Term::Error.delay())
+                .apply(
+                    Term::sha3_256()
+                        .apply(Term::sha3_256().apply(Term::byte_string(vec![])))
+                        .delay(),
+                )
+                .force()
+                .lambda("__if_then_else_wrapped")
+                .apply(Term::Builtin(DefaultFunction::IfThenElse).force()),
+        };
+
+        compare_optimization(expected, program, |p| {
+            p.run_one_opt(true, &mut |id, term, arg_stack, scope, context| {
+                term.inline_reducer(id, arg_stack, scope, context);
+            })
+        });
+    }
+
     #[test]
     fn inline_reduce_0_occurrence() {
         let program: Program<Name> = Program {