From 5e2a78173bfd4a24ed6746e3f4f94495ce06d93e Mon Sep 17 00:00:00 2001
From: microproofs <kwhitemsg@gmail.com>
Date: Fri, 29 Sep 2023 14:26:55 -0400
Subject: [PATCH] optimization: increase the cases that inline_basic_reduce can
 handle

---
 crates/uplc/src/optimize.rs          |  2 ++
 crates/uplc/src/optimize/shrinker.rs | 24 ++++++++++++++++++++----
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/crates/uplc/src/optimize.rs b/crates/uplc/src/optimize.rs
index e660d1d5..02501b5a 100644
--- a/crates/uplc/src/optimize.rs
+++ b/crates/uplc/src/optimize.rs
@@ -26,4 +26,6 @@ pub fn aiken_optimize_and_intern(program: Program<Name>) -> Program<Name> {
         .wrap_data_reduce()
         .lambda_reduce()
         .inline_reduce()
+        .lambda_reduce()
+        .inline_reduce()
 }
diff --git a/crates/uplc/src/optimize/shrinker.rs b/crates/uplc/src/optimize/shrinker.rs
index 8191bde1..efe0b6b0 100644
--- a/crates/uplc/src/optimize/shrinker.rs
+++ b/crates/uplc/src/optimize/shrinker.rs
@@ -283,17 +283,21 @@ fn inline_basic_reduce(term: &mut Term<Name>) {
             } = func
             {
                 let occurrences = var_occurrences(body, parameter_name.clone());
+                // TODO: Once case and constr is live we need to also check if case is delaying values
+                let delays = delayed_execution(body.as_ref());
                 if occurrences == 1 {
-                    if let replace_term @ (Term::Var(_)
+                    if delays == 0 {
+                        *term = substitute_term(body.as_ref(), parameter_name.clone(), arg);
+                    } else if let Term::Var(_)
                     | Term::Constant(_)
                     | Term::Error
                     | Term::Delay(_)
                     | Term::Lambda { .. }
-                    | Term::Builtin(_)) = arg
+                    | Term::Builtin(_) = arg
                     {
-                        *term =
-                            substitute_term(body.as_ref(), parameter_name.clone(), replace_term);
+                        *term = substitute_term(body.as_ref(), parameter_name.clone(), arg);
                     }
+                // This will strip out unused terms that can't throw an error by themselves
                 } else if occurrences == 0 {
                     if let Term::Var(_)
                     | Term::Constant(_)
@@ -464,6 +468,18 @@ fn var_occurrences(term: &Term<Name>, search_for: Rc<Name>) -> usize {
     }
 }
 
+fn delayed_execution(term: &Term<Name>) -> usize {
+    match term {
+        Term::Delay(body) => 1 + delayed_execution(body.as_ref()),
+        Term::Lambda { body, .. } => 1 + delayed_execution(body.as_ref()),
+        Term::Apply { function, argument } => {
+            delayed_execution(function.as_ref()) + delayed_execution(argument.as_ref())
+        }
+        Term::Force(x) => delayed_execution(x.as_ref()),
+        _ => 0,
+    }
+}
+
 fn lambda_reduce(term: &mut Term<Name>) {
     match term {
         Term::Apply { function, argument } => {