Handle case where write_bits is used without being applied

aiken-lang · Jan 11, 2025 · 90fa7f6 · 90fa7f6
1 parent d9533e8
commit 90fa7f6
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 3 deletions.
diff --git a/crates/uplc/src/builder.rs b/crates/uplc/src/builder.rs
@@ -408,6 +408,10 @@ where
     pub fn serialise_data() -> Self {
         Term::Builtin(DefaultFunction::SerialiseData)
     }
+
+    pub fn write_bits() -> Self {
+        Term::Builtin(DefaultFunction::WriteBits)
+    }
 }
 
 impl<T> Term<T>

diff --git a/crates/uplc/src/optimize/shrinker.rs b/crates/uplc/src/optimize/shrinker.rs
@@ -1498,6 +1498,7 @@ impl Term<Name> {
                         }
                         arg => {
                             context.write_bits_convert = true;
+
                             *arg = Term::var(INDICES_CONVERTER)
                                 .apply(std::mem::replace(arg, Term::Error.force()));
                         }
@@ -1506,10 +1507,24 @@ impl Term<Name> {
             }
 
             Term::Builtin(DefaultFunction::WriteBits) => {
-                // first arg not needed
-                arg_stack.pop();
+                if arg_stack.is_empty() {
+                    context.write_bits_convert = true;
+
+                    *self = Term::write_bits()
+                        .apply(Term::var("__arg_1"))
+                        .apply(Term::var(INDICES_CONVERTER).apply(Term::var("__arg_2")))
+                        .apply(Term::var("__arg_3"))
+                        .lambda("__arg_3")
+                        .lambda("__arg_2")
+                        .lambda("__arg_1")
+                } else {
+                    // first arg not needed
+                    arg_stack.pop();
+
+                    let Some(Args::Apply(arg_id, _)) = arg_stack.pop() else {
+                        return;
+                    };
 
-                if let Some(Args::Apply(arg_id, _)) = arg_stack.pop() {
                     context.write_bits_indices_arg.push(arg_id);
                 }
             }

diff --git a/examples/acceptance_tests/117/lib/tests.ak b/examples/acceptance_tests/117/lib/tests.ak
@@ -15,3 +15,22 @@ test baz() {
   let x = [0, 1, 2, 3]
   write_bits(#"f0", x, True) == #"ff"
 }
+
+test bur() {
+  let x =
+    if True {
+      [0, 1, 2, 3]
+    } else {
+      [0, 1]
+    }
+
+  if False {
+    fn(_a, _b, _c) { #"" }
+  } else {
+    write_bits
+  }(
+    #"f0",
+    x,
+    True,
+  ) == #"ff"
+}