Allow write_bits builtin and add a few more optimizations

crates/aiken-lang/src/builtins.rs

@@ -509,12 +509,8 @@ pub fn plutus(id_gen: &IdGenerator) -> TypeInfo {
     };
 
     for builtin in DefaultFunction::iter() {
-        // FIXME: Disabling WriteBits for now, since its signature requires the ability to create
-        // list of raw integers, which isn't possible through Aiken at the moment.
-        if !matches!(builtin, DefaultFunction::WriteBits) {
-            let value = from_default_function(builtin, id_gen);
-            plutus.values.insert(builtin.aiken_name(), value);
-        }
+        let value = from_default_function(builtin, id_gen);
+        plutus.values.insert(builtin.aiken_name(), value);
     }
 
     let index_tipo = Type::function(vec![Type::data()], Type::int());

crates/aiken-lang/src/gen_uplc.rs

@@ -3712,7 +3712,7 @@ impl<'a> CodeGenerator<'a> {
                     interner.program(&mut program);
 
                     let eval_program: Program<NamedDeBruijn> =
-                        program.clean_up().try_into().unwrap();
+                        program.clean_up(false).try_into().unwrap();
 
                     Some(
                         eval_program
@@ -3822,7 +3822,7 @@ impl<'a> CodeGenerator<'a> {
                             interner.program(&mut program);
 
                             let eval_program: Program<NamedDeBruijn> =
-                                program.clean_up().try_into().unwrap();
+                                program.clean_up(false).try_into().unwrap();
 
                             let evaluated_term: Term<NamedDeBruijn> = eval_program
                                 .eval(ExBudget::default())
@@ -4028,7 +4028,7 @@ impl<'a> CodeGenerator<'a> {
                 } else {
                     let term = arg_stack.pop().unwrap();
 
-                    match term.pierce_no_inlines() {
+                    match term.pierce_no_inlines_ref() {
                         Term::Var(_) => Some(term.force()),
                         Term::Delay(inner_term) => Some(inner_term.as_ref().clone()),
                         Term::Apply { .. } => Some(term.force()),
@@ -4356,15 +4356,15 @@ impl<'a> CodeGenerator<'a> {
                     known_data_to_type(term, &tipo)
                 };
 
-                if extract_constant(term.pierce_no_inlines()).is_some() {
+                if extract_constant(term.pierce_no_inlines_ref()).is_some() {
                     let mut program = self.new_program(term);
 
                     let mut interner = CodeGenInterner::new();
 
                     interner.program(&mut program);
 
                     let eval_program: Program<NamedDeBruijn> =
-                        program.clean_up().try_into().unwrap();
+                        program.clean_up(false).try_into().unwrap();
 
                     let evaluated_term: Term<NamedDeBruijn> = eval_program
                         .eval(ExBudget::default())
@@ -4379,7 +4379,7 @@ impl<'a> CodeGenerator<'a> {
             Air::CastToData { tipo } => {
                 let mut term = arg_stack.pop().unwrap();
 
-                if extract_constant(term.pierce_no_inlines()).is_some() {
+                if extract_constant(term.pierce_no_inlines_ref()).is_some() {
                     term = builder::convert_type_to_data(term, &tipo);
 
                     let mut program = self.new_program(term);
@@ -4389,7 +4389,7 @@ impl<'a> CodeGenerator<'a> {
                     interner.program(&mut program);
 
                     let eval_program: Program<NamedDeBruijn> =
-                        program.clean_up().try_into().unwrap();
+                        program.clean_up(false).try_into().unwrap();
 
                     let evaluated_term: Term<NamedDeBruijn> = eval_program
                         .eval(ExBudget::default())
@@ -4792,7 +4792,7 @@ impl<'a> CodeGenerator<'a> {
                     .apply(term);
 
                 if arg_vec.iter().all(|item| {
-                    let maybe_const = extract_constant(item.pierce_no_inlines());
+                    let maybe_const = extract_constant(item.pierce_no_inlines_ref());
                     maybe_const.is_some()
                 }) {
                     let mut program = self.new_program(term);
@@ -4802,7 +4802,7 @@ impl<'a> CodeGenerator<'a> {
                     interner.program(&mut program);
 
                     let eval_program: Program<NamedDeBruijn> =
-                        program.clean_up().try_into().unwrap();
+                        program.clean_up(false).try_into().unwrap();
 
                     let evaluated_term: Term<NamedDeBruijn> = eval_program
                         .eval(ExBudget::default())

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",
-  "hash": "c6af3f04e300cb8c1d0429cc0d8e56a0413eef9fcb338f72076b426c",
+  "compiledCode": "5901870101009800aba2aba1aba0aab9eaab9dab9a488888888c8c8c8c8c966002600860126ea8006264b30013005300a375400314800226466e00dd698070009980226103d8798000300e300f001300b37540028048c030c034016264b30013370e900118051baa0018991919b80337006eb4c03c008dd6980780099802980798080011807980800098061baa002300b3754005132337006eb4c038004cc010c038c03c00530103d8798000300b37540048048c030c0340150081805802180080091119192cc004c018c02cdd5000c4c966002600e60186ea80062900044c8cdc01bad30100019800803d300103d879800098081808800a00e300d37540028058c038c03c00a264b30013370e900118061baa0018991919b80337006eb4c044008dd69808800cc00402260226024005301130120014020601c6ea8008c034dd500144c8cdc01bad30100019800803cc040c044006980103d8798000401c601a6ea800900b180718078012014300d0013300b0023300b0014bd701b8748000cc018008cc0180052f5c01",
+  "hash": "247535960781372d3b2097595ebd748bd61be7c8f2f264e460e095b3",
   "definitions": {
     "Int": {
       "dataType": "integer"

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

@@ -3603,7 +3603,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 +3627,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 +3651,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,16 +4088,16 @@ 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()
+                                Term::tail_list()
                                     .apply(Term::Var(tail_id_6.clone()))
-                                    .as_var("__val", |val| {
-                                        Term::tail_list()
+                                    .delayed_choose_list(
+                                        Term::head_list()
                                             .apply(Term::Var(tail_id_6))
-                                            .delayed_choose_list(
-                                                expect_b(val, Term::Var(then_delayed), trace),
-                                                Term::Error,
-                                            )
-                                    })
+                                            .as_var("__val", |val| {
+                                                expect_b(val, Term::Var(then_delayed), trace)
+                                            }),
+                                        Term::Error,
+                                    )
                             })
                     }
                 });

crates/uplc/src/ast.rs

@@ -517,7 +517,7 @@ impl hash::Hash for Name {
 
 impl PartialEq for Name {
     fn eq(&self, other: &Self) -> bool {
-        self.unique == other.unique
+        self.unique == other.unique && self.text == other.text
     }
 }
 

crates/uplc/src/builder.rs

@@ -9,6 +9,7 @@ pub const CONSTR_FIELDS_EXPOSER: &str = "__constr_fields_exposer";
 pub const CONSTR_INDEX_EXPOSER: &str = "__constr_index_exposer";
 pub const EXPECT_ON_LIST: &str = "__expect_on_list";
 pub const INNER_EXPECT_ON_LIST: &str = "__inner_expect_on_list";
+pub const INDICES_CONVERTER: &str = "__indices_converter";
 
 impl<T> Term<T>
 where
@@ -30,6 +31,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())
@@ -71,6 +83,10 @@ where
         Term::Constant(Constant::ProtoList(Type::Data, vals).into())
     }
 
+    pub fn int_values(vals: Vec<Constant>) -> Self {
+        Term::Constant(Constant::ProtoList(Type::Integer, vals).into())
+    }
+
     pub fn empty_map() -> Self {
         Term::Constant(
             Constant::ProtoList(Type::Pair(Type::Data.into(), Type::Data.into()), vec![]).into(),
@@ -392,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>
@@ -535,6 +555,33 @@ impl Term<Name> {
         )
     }
 
+    pub fn data_list_to_integer_list(self) -> Self {
+        self.lambda(INDICES_CONVERTER)
+            .apply(Term::var(INDICES_CONVERTER).apply(Term::var(INDICES_CONVERTER)))
+            .lambda(INDICES_CONVERTER)
+            .apply(
+                Term::var("xs")
+                    .delayed_choose_list(
+                        Term::int_values(vec![]),
+                        Term::mk_cons()
+                            .apply(Term::var("x"))
+                            .apply(
+                                Term::var(INDICES_CONVERTER)
+                                    .apply(Term::var(INDICES_CONVERTER))
+                                    .apply(Term::var("rest")),
+                            )
+                            .lambda("rest")
+                            .apply(Term::tail_list().apply(Term::var("xs")))
+                            .lambda("x")
+                            .apply(
+                                Term::un_i_data().apply(Term::head_list().apply(Term::var("xs"))),
+                            ),
+                    )
+                    .lambda("xs")
+                    .lambda(INDICES_CONVERTER),
+            )
+    }
+
     /// Introduce a let-binding for a given term. The callback receives a Term::Var
     /// whose name matches the given 'var_name'. Handy to re-use a same var across
     /// multiple lambda expressions.

crates/uplc/src/builtins.rs

@@ -6,7 +6,19 @@ use strum_macros::EnumIter;
 /// All the possible builtin functions in Untyped Plutus Core.
 #[repr(u8)]
 #[allow(non_camel_case_types)]
-#[derive(Debug, Clone, PartialEq, Eq, Copy, EnumIter, serde::Serialize, serde::Deserialize)]
+#[derive(
+    Debug,
+    Clone,
+    PartialEq,
+    Eq,
+    Copy,
+    EnumIter,
+    serde::Serialize,
+    serde::Deserialize,
+    Hash,
+    PartialOrd,
+    Ord,
+)]
 pub enum DefaultFunction {
     // Integer functions
     AddInteger = 0,

crates/uplc/src/optimize.rs

@@ -38,5 +38,5 @@ pub fn aiken_optimize_and_intern(program: Program<Name>) -> Program<Name> {
         }
     }
 
-    prog.clean_up()
+    prog.clean_up(true)
 }