expect on map and list now use a function that gets instantiated as air in code gen functions

crates/aiken-lang/src/gen_uplc.rs

@@ -6,7 +6,7 @@ use uplc::{
     ast::{
         Constant as UplcConstant, DeBruijn, Name, NamedDeBruijn, Program, Term, Type as UplcType,
     },
-    builder::{CONSTR_FIELDS_EXPOSER, CONSTR_GET_FIELD, CONSTR_INDEX_EXPOSER},
+    builder::{CONSTR_FIELDS_EXPOSER, CONSTR_GET_FIELD, CONSTR_INDEX_EXPOSER, EXPECT_ON_LIST},
     builtins::DefaultFunction,
     machine::cost_model::ExBudget,
     optimize::aiken_optimize_and_intern,
@@ -18,7 +18,7 @@ use crate::{
         ArgName, AssignmentKind, BinOp, Pattern, Span, TypedArg, TypedClause, TypedDataType,
         TypedFunction, TypedValidator, UnOp,
     },
-    builtins::{bool, data},
+    builtins::{bool, data, void},
     expr::TypedExpr,
     tipo::{
         ModuleValueConstructor, PatternConstructor, Type, TypeInfo, ValueConstructor,
@@ -40,10 +40,9 @@ use builder::{
 use self::{builder::replace_opaque_type, scope::Scope, stack::AirStack};
 
 #[derive(Clone, Debug)]
-pub struct CodeGenFunction {
-    air: Vec<Air>,
-    dependencies: Vec<String>,
-    recursive: bool,
+pub enum CodeGenFunction {
+    Function(Vec<Air>, Vec<String>),
+    Link(String),
 }
 
 #[derive(Clone)]
@@ -147,7 +146,7 @@ impl<'a> CodeGenerator<'a> {
 
         term = self.wrap_validator_args(term, params, false);
 
-        self.finalize(term, true)
+        self.finalize(term)
     }
 
     pub fn generate_test(&mut self, test_body: &TypedExpr) -> Program<Name> {
@@ -165,10 +164,10 @@ impl<'a> CodeGenerator<'a> {
 
         let term = self.uplc_code_gen(&mut ir_stack);
 
-        self.finalize(term, false)
+        self.finalize(term)
     }
 
-    fn finalize(&mut self, term: Term<Name>, wrap_as_validator: bool) -> Program<Name> {
+    fn finalize(&mut self, term: Term<Name>) -> Program<Name> {
         // let mut term = if self.used_data_assert_on_list {
         //     term.assert_on_list()
         // } else {
@@ -2274,6 +2273,21 @@ impl<'a> CodeGenerator<'a> {
             unwrap_function_stack
                 .anonymous_function(vec![format!("__pair_{new_id}")], pair_access_stack);
 
+            let function = self.code_gen_functions.get(EXPECT_ON_LIST);
+
+            if function.is_none() {
+                let mut expect_list_stack = expect_stack.empty_with_scope();
+
+                expect_list_stack.expect_on_list();
+                self.code_gen_functions.insert(
+                    EXPECT_ON_LIST.to_string(),
+                    CodeGenFunction::Function(
+                        expect_list_stack.complete(),
+                        vec!["__list_to_check".to_string(), "__check_with".to_string()],
+                    ),
+                );
+            }
+
             expect_stack.expect_list_from_data(tipo.clone(), name, unwrap_function_stack);
 
             expect_stack.void();
@@ -2300,6 +2314,21 @@ impl<'a> CodeGenerator<'a> {
             unwrap_function_stack
                 .anonymous_function(vec![format!("__list_item_{new_id}")], list_access_stack);
 
+            let function = self.code_gen_functions.get(EXPECT_ON_LIST);
+
+            if function.is_none() {
+                let mut expect_list_stack = expect_stack.empty_with_scope();
+
+                expect_list_stack.expect_on_list();
+                self.code_gen_functions.insert(
+                    EXPECT_ON_LIST.to_string(),
+                    CodeGenFunction::Function(
+                        expect_list_stack.complete(),
+                        vec!["__list_to_check".to_string(), "__check_with".to_string()],
+                    ),
+                );
+            }
+
             expect_stack.expect_list_from_data(tipo.clone(), name, unwrap_function_stack);
 
             expect_stack.void();

crates/aiken-lang/src/gen_uplc/stack.rs

@@ -269,7 +269,21 @@ impl AirStack {
             tipo: tipo.clone(),
         });
 
-        self.local_var(tipo.clone(), EXPECT_ON_LIST);
+        self.var(
+            ValueConstructor::public(
+                void(),
+                ValueConstructorVariant::ModuleFn {
+                    name: EXPECT_ON_LIST.to_string(),
+                    field_map: None,
+                    module: "".to_string(),
+                    arity: 2,
+                    location: Span::empty(),
+                    builtin: None,
+                },
+            ),
+            EXPECT_ON_LIST,
+            "",
+        );
 
         self.local_var(tipo, name);
 
@@ -707,7 +721,7 @@ impl AirStack {
         self.merge_child(value_stack);
     }
 
-    pub fn assert_on_list(&mut self) {
+    pub fn expect_on_list(&mut self) {
         let mut head_stack = self.empty_with_scope();
         let mut tail_stack = self.empty_with_scope();
         let mut check_with_stack = self.empty_with_scope();
@@ -740,7 +754,7 @@ impl AirStack {
 
         self.list_clause(
             void(),
-            "__list_to_check".to_string(),
+            "__list_to_check",
             None,
             false,
             void_stack,
@@ -748,11 +762,25 @@ impl AirStack {
 
         self.choose_unit(check_with_stack);
 
-        expect_stack.local_var(void(), EXPECT_ON_LIST.to_string());
+        expect_stack.var(
+            ValueConstructor::public(
+                void(),
+                ValueConstructorVariant::ModuleFn {
+                    name: EXPECT_ON_LIST.to_string(),
+                    field_map: None,
+                    module: "".to_string(),
+                    arity: 2,
+                    location: Span::empty(),
+                    builtin: None,
+                },
+            ),
+            EXPECT_ON_LIST,
+            "",
+        );
 
-        arg_stack1.local_var(list(data()), "__list_to_check".to_string());
+        arg_stack1.local_var(list(data()), "__list_to_check");
 
-        arg_stack2.local_var(void(), "__check_with".to_string());
+        arg_stack2.local_var(void(), "__check_with");
 
         tail_stack.builtin(DefaultFunction::TailList, list(data()), vec![arg_stack1]);
 
@@ -885,7 +913,7 @@ mod test {
             air: vec![],
         };
 
-        stack1.assert_on_list();
+        stack1.expect_on_list();
 
         println!("{:#?}", stack1);
 
@@ -959,13 +987,17 @@ mod test {
             },
             Air::Var {
                 scope: vec![0, 7, 8, 13, 9].into(),
-                constructor: ValueConstructor {
-                    public: true,
-                    variant: ValueConstructorVariant::LocalVariable {
+                constructor: ValueConstructor::public(
+                    void(),
+                    ValueConstructorVariant::ModuleFn {
+                        name: "__expect_on_list".to_string(),
+                        field_map: None,
+                        module: "".to_string(),
+                        arity: 2,
                         location: Span::empty(),
+                        builtin: None,
                     },
-                    tipo: void(),
-                },
+                ),
                 name: "__expect_on_list".to_string(),
                 variant_name: "".to_string(),
             },