Update script context handler to do less work with only fallback. Also optimize wrapped builtins too.

crates/aiken-lang/src/ast.rs

@@ -567,7 +567,7 @@ impl TypedValidator {
         let var_purpose_arg = "__purpose_arg__";
         let var_datum = "__datum__";
 
-        TypedExpr::sequence(&[
+        let context_handler = TypedExpr::sequence(&[
             TypedExpr::let_(
                 TypedExpr::local_var(var_context, Type::script_context(), self.location),
                 TypedPattern::Constructor {
@@ -746,7 +746,29 @@ impl TypedValidator {
                     }))
                     .collect(),
             },
-        ])
+        ]);
+
+        if self.handlers.is_empty() {
+            let fallback = &self.fallback;
+            let arg = fallback.arguments.first().unwrap();
+
+            let then = match arg.get_variable_name() {
+                Some(arg_name) => TypedExpr::sequence(&[
+                    TypedExpr::let_(
+                        TypedExpr::local_var(var_context, arg.tipo.clone(), arg.location),
+                        TypedPattern::var(arg_name),
+                        arg.tipo.clone(),
+                        arg.location,
+                    ),
+                    fallback.body.clone(),
+                ]),
+                None => fallback.body.clone(),
+            };
+
+            then
+        } else {
+            context_handler
+        }
     }
 
     pub fn find_node(&self, byte_index: usize) -> Option<Located<'_>> {

crates/uplc/src/optimize/shrinker.rs

@@ -1316,7 +1316,7 @@ impl Term<Name> {
                         // So it costs more size to have them hoisted
                         Term::Delay(e) if matches!(e.as_ref(), Term::Error) => true,
                         // If it wraps a builtin with consts or arguments passed in then inline
-                        l @ Term::Lambda { .. } if is_a_builtin_wrapper(l) => true,
+                        Term::Lambda { .. } => arg_term.is_a_builtin_wrapper(&context),
                         // Inline smaller terms too
                         Term::Constant(_) | Term::Var(_) | Term::Builtin(_) => true,
 
@@ -2194,7 +2194,7 @@ impl Term<Name> {
         term
     }
 
-    pub fn pierce_no_inlines(mut self) -> Self {
+    fn pierce_no_inlines(mut self) -> Self {
         let term = &mut self;
 
         while let Term::Lambda {
@@ -2211,6 +2211,59 @@ impl Term<Name> {
 
         std::mem::replace(term, Term::Error.force())
     }
+
+    fn is_a_builtin_wrapper(&self, context: &Context) -> bool {
+        let (names, term) = self.pop_lambdas_and_get_names();
+
+        let mut arg_names = vec![];
+
+        let mut term = term;
+
+        while let Term::Apply { function, argument } = term {
+            match argument.as_ref() {
+                Term::Var(name) => arg_names.push(format!("{}_{}", name.text, name.unique)),
+
+                Term::Constant(_) => {}
+                _ => {
+                    //Break loop, it's not a builtin wrapper function
+                    return false;
+                }
+            }
+            term = function.as_ref();
+        }
+
+        let func_is_builtin = match term {
+            Term::Var(name) => context
+                .builtins_map
+                .keys()
+                .map(|func| func.wrapped_name())
+                .any(|func| func == name.text),
+
+            Term::Builtin(_) => todo!(),
+            _ => false,
+        };
+
+        arg_names.iter().all(|item| names.contains(item)) && func_is_builtin
+    }
+
+    fn pop_lambdas_and_get_names(&self) -> (Vec<String>, &Term<Name>) {
+        let mut names = vec![];
+
+        let mut term = self;
+
+        while let Term::Lambda {
+            parameter_name,
+            body,
+        } = term
+        {
+            if parameter_name.text != NO_INLINE {
+                names.push(format!("{}_{}", parameter_name.text, parameter_name.unique));
+            }
+            term = body.as_ref();
+        }
+
+        (names, term)
+    }
 }
 
 impl Program<Name> {
@@ -2651,48 +2704,6 @@ fn id_vec_function_to_var(func_name: &str, id_vec: &[usize]) -> String {
     )
 }
 
-fn is_a_builtin_wrapper(term: &Term<Name>) -> bool {
-    let (names, term) = pop_lambdas_and_get_names(term);
-
-    let mut arg_names = vec![];
-
-    let mut term = term;
-
-    while let Term::Apply { function, argument } = term {
-        match argument.as_ref() {
-            Term::Var(name) => arg_names.push(format!("{}_{}", name.text, name.unique)),
-
-            Term::Constant(_) => {}
-            _ => {
-                //Break loop, it's not a builtin wrapper function
-                return false;
-            }
-        }
-        term = function.as_ref();
-    }
-
-    arg_names.iter().all(|item| names.contains(item)) && matches!(term, Term::Builtin(_))
-}
-
-fn pop_lambdas_and_get_names(term: &Term<Name>) -> (Vec<String>, &Term<Name>) {
-    let mut names = vec![];
-
-    let mut term = term;
-
-    while let Term::Lambda {
-        parameter_name,
-        body,
-    } = term
-    {
-        if parameter_name.text != NO_INLINE {
-            names.push(format!("{}_{}", parameter_name.text, parameter_name.unique));
-        }
-        term = body.as_ref();
-    }
-
-    (names, term)
-}
-
 #[cfg(test)]
 mod tests {
     use super::NO_INLINE;