Pushing up what I have so far

aiken-lang · Aug 30, 2024 · 28fd2f1 · 28fd2f1
1 parent 958caa8
commit 28fd2f1
Show file tree

Hide file tree

Showing 3 changed files with 146 additions and 37 deletions.
diff --git a/crates/aiken-lang/src/gen_uplc.rs b/crates/aiken-lang/src/gen_uplc.rs
@@ -3623,40 +3623,51 @@ impl<'a> CodeGenerator<'a> {
 
             sorting_attempts += 1;
         }
-        sorted_function_vec.dedup();
 
         let mut sorted_hoistable = constants_to_hoist
             .keys()
             .map(|k| (k.to_owned(), "".to_string()))
             .collect::<Vec<_>>();
         sorted_hoistable.extend(sorted_function_vec);
 
+        sorted_hoistable.dedup();
+
         // Now we need to hoist the functions to the top of the validator
         for (key, variant) in sorted_hoistable {
-            if hoisted_functions
-                .iter()
-                .any(|(func_key, func_variant)| func_key == &key && func_variant == &variant)
-            {
-                continue;
-            }
+            if let Some((tree_path, value)) = constants_to_hoist.get(&key) {
+                self.hoist_function(
+                    &mut air_tree,
+                    tree_path,
+                    value,
+                    (&key, &variant),
+                    &functions_to_hoist,
+                    &mut hoisted_functions,
+                );
+            } else {
+                if hoisted_functions
+                    .iter()
+                    .any(|(func_key, func_variant)| func_key == &key && func_variant == &variant)
+                {
+                    continue;
+                }
 
-            let function_variants = functions_to_hoist
-                .get(&key)
-                .unwrap_or_else(|| panic!("Missing Function Definition"));
+                let function_variants = functions_to_hoist
+                    .get(&key)
+                    .unwrap_or_else(|| panic!("Missing Function Definition"));
 
-            let (tree_path, function) = function_variants
-                .get(&variant)
-                .unwrap_or_else(|| panic!("Missing Function Variant Definition"));
+                let (tree_path, function) = function_variants
+                    .get(&variant)
+                    .unwrap_or_else(|| panic!("Missing Function Variant Definition"));
 
-            self.hoist_function(
-                &mut air_tree,
-                tree_path,
-                function,
-                (&key, &variant),
-                &functions_to_hoist,
-                &constants_to_hoist,
-                &mut hoisted_functions,
-            );
+                self.hoist_function(
+                    &mut air_tree,
+                    tree_path,
+                    function,
+                    (&key, &variant),
+                    &functions_to_hoist,
+                    &mut hoisted_functions,
+                );
+            }
         }
 
         air_tree
@@ -3669,11 +3680,23 @@ impl<'a> CodeGenerator<'a> {
         hoistable: &Hoistable,
         key_var: (&FunctionAccessKey, &String),
         functions_to_hoist: &IndexMap<FunctionAccessKey, IndexMap<String, (TreePath, Hoistable)>>,
-        constants_to_hoist: &IndexMap<FunctionAccessKey, (TreePath, Hoistable)>,
         hoisted_functions: &mut Vec<(FunctionAccessKey, String)>,
     ) {
         match hoistable {
-            Hoistable::Constant { .. } => {}
+            Hoistable::Constant { value } => {
+                let access_key = key_var.0;
+
+                let node_to_edit = air_tree.find_air_tree_node(tree_path);
+
+                let defined_const = AirTree::define_hoisted_constant(
+                    &access_key.function_name,
+                    &access_key.module_name,
+                    value.clone(),
+                    std::mem::replace(node_to_edit, AirTree::void()),
+                );
+
+                *node_to_edit = defined_const;
+            }
             Hoistable::Function {
                 body,
                 deps: func_deps,
@@ -3709,7 +3732,7 @@ impl<'a> CodeGenerator<'a> {
                     is_recursive,
                     recursive_nonstatics,
                     body,
-                    node_to_edit.clone(),
+                    std::mem::replace(node_to_edit, AirTree::void()),
                 );
 
                 let defined_dependencies = self.hoist_dependent_functions(
@@ -3746,7 +3769,7 @@ impl<'a> CodeGenerator<'a> {
                     &key.module_name,
                     variant,
                     functions,
-                    node_to_edit.clone(),
+                    std::mem::replace(node_to_edit, AirTree::void()),
                 );
 
                 let defined_dependencies = self.hoist_dependent_functions(
@@ -3792,7 +3815,7 @@ impl<'a> CodeGenerator<'a> {
                 .unwrap_or_else(|| panic!("Missing Function Variant Definition"));
 
             match function {
-                Hoistable::Constant { .. } => todo!("Hoistable constant handler."),
+                Hoistable::Constant { .. } => unreachable!(),
                 Hoistable::Function { deps, .. } => {
                     for (dep_generic_func, dep_variant) in deps.iter() {
                         if !(dep_generic_func == &dep.0 && dep_variant == &dep.1) {
@@ -4278,13 +4301,22 @@ impl<'a> CodeGenerator<'a> {
                     }
                     .into(),
                 )),
-                ValueConstructorVariant::ModuleConstant { module, name, .. } => Some(Term::Var(
-                    Name {
-                        text: format!("{module}_{name}"),
-                        unique: 0.into(),
-                    }
-                    .into(),
-                )),
+
+                ValueConstructorVariant::ModuleConstant { module, name, .. } => {
+                    let name = if module.is_empty() {
+                        name.to_string()
+                    } else {
+                        format!("{module}_{name}")
+                    };
+                    Some(Term::Var(
+                        Name {
+                            text: name,
+                            unique: 0.into(),
+                        }
+                        .into(),
+                    ))
+                }
+
                 ValueConstructorVariant::ModuleFn {
                     name: func_name,
                     module,
@@ -4856,6 +4888,21 @@ impl<'a> CodeGenerator<'a> {
                     Some(term)
                 }
             }
+            Air::DefineConstant {
+                const_name,
+                module_name,
+                value,
+            } => {
+                let name = if module_name.is_empty() {
+                    const_name.to_string()
+                } else {
+                    format!("{module_name}_{const_name}")
+                };
+
+                let term = arg_stack.pop().unwrap();
+
+                Some(term.lambda(name).apply(value))
+            }
             Air::DefineCyclicFuncs {
                 func_name,
                 module_name,

diff --git a/crates/aiken-lang/src/gen_uplc/air.rs b/crates/aiken-lang/src/gen_uplc/air.rs
@@ -4,7 +4,10 @@ use crate::{
 };
 use indexmap::IndexSet;
 use std::rc::Rc;
-use uplc::builtins::DefaultFunction;
+use uplc::{
+    ast::{Name, Term},
+    builtins::DefaultFunction,
+};
 
 #[derive(Debug, Clone, PartialEq, Copy)]
 pub enum ExpectLevel {
@@ -72,6 +75,11 @@ pub enum Air {
         recursive_nonstatic_params: Vec<String>,
         variant_name: String,
     },
+    DefineConstant {
+        const_name: String,
+        module_name: String,
+        value: Term<Name>,
+    },
     DefineCyclicFuncs {
         func_name: String,
         module_name: String,

diff --git a/crates/aiken-lang/src/gen_uplc/tree.rs b/crates/aiken-lang/src/gen_uplc/tree.rs
@@ -7,6 +7,7 @@ use indexmap::IndexSet;
 use itertools::Itertools;
 use std::{borrow::BorrowMut, rc::Rc, slice::Iter};
 use uplc::{
+    ast::{Name, Term},
     builder::{EXPECT_ON_LIST, INNER_EXPECT_ON_LIST},
     builtins::DefaultFunction,
 };
@@ -143,6 +144,12 @@ pub enum AirTree {
         func_body: Box<AirTree>,
         then: Box<AirTree>,
     },
+    DefineConstant {
+        const_name: String,
+        module_name: String,
+        value: Term<Name>,
+        then: Box<AirTree>,
+    },
     DefineCyclicFuncs {
         func_name: String,
         module_name: String,
@@ -537,6 +544,20 @@ impl AirTree {
         }
     }
 
+    pub fn define_hoisted_constant(
+        const_name: impl ToString,
+        module_name: impl ToString,
+        value: Term<Name>,
+        then: AirTree,
+    ) -> AirTree {
+        AirTree::DefineConstant {
+            const_name: const_name.to_string(),
+            module_name: module_name.to_string(),
+            value,
+            then: then.into(),
+        }
+    }
+
     pub fn define_cyclic_func(
         func_name: impl ToString,
         module_name: impl ToString,
@@ -1173,6 +1194,19 @@ impl AirTree {
                 func_body.create_air_vec(air_vec);
                 then.create_air_vec(air_vec);
             }
+            AirTree::DefineConstant {
+                const_name,
+                module_name,
+                value,
+                then,
+            } => {
+                air_vec.push(Air::DefineConstant {
+                    const_name: const_name.clone(),
+                    module_name: module_name.clone(),
+                    value: value.clone(),
+                });
+                then.create_air_vec(air_vec);
+            }
             AirTree::DefineCyclicFuncs {
                 func_name,
                 module_name,
@@ -1712,6 +1746,7 @@ impl AirTree {
             | AirTree::Let { then, .. }
             | AirTree::SoftCastLet { then, .. }
             | AirTree::DefineFunc { then, .. }
+            | AirTree::DefineConstant { then, .. }
             | AirTree::DefineCyclicFuncs { then, .. }
             | AirTree::AssertConstr { then, .. }
             | AirTree::AssertBool { then, .. }
@@ -1790,6 +1825,7 @@ impl AirTree {
             }
             AirTree::Let { .. }
             | AirTree::DefineFunc { .. }
+            | AirTree::DefineConstant { .. }
             | AirTree::DefineCyclicFuncs { .. }
             | AirTree::AssertConstr { .. }
             | AirTree::AssertBool { .. }
@@ -2129,6 +2165,7 @@ impl AirTree {
             }
 
             AirTree::DefineFunc { .. }
+            | AirTree::DefineConstant { .. }
             | AirTree::DefineCyclicFuncs { .. }
             | AirTree::ListClauseGuard { .. }
             | AirTree::TupleGuard { .. }
@@ -2580,6 +2617,18 @@ impl AirTree {
                     apply_with_func_last,
                 )
             }
+            AirTree::DefineConstant {
+                const_name: _,
+                module_name: _,
+                value: _,
+                then,
+            } => then.do_traverse_tree_with(
+                tree_path,
+                current_depth + 1,
+                Fields::FourthField,
+                with,
+                apply_with_func_last,
+            ),
             AirTree::DefineCyclicFuncs {
                 func_name: _,
                 module_name: _,
@@ -2873,7 +2922,10 @@ impl AirTree {
         tree_path_iter: &mut Iter<(usize, Fields)>,
     ) -> &'a mut AirTree {
         // For finding the air node we skip over the define func ops since those are added later on.
-        if let AirTree::DefineFunc { then, .. } | AirTree::DefineCyclicFuncs { then, .. } = self {
+        if let AirTree::DefineFunc { then, .. }
+        | AirTree::DefineCyclicFuncs { then, .. }
+        | AirTree::DefineConstant { then, .. } = self
+        {
             then.as_mut().do_find_air_tree_node(tree_path_iter)
         } else if let Some((_depth, field)) = tree_path_iter.next() {
             match self {
@@ -3029,7 +3081,9 @@ impl AirTree {
                     Fields::FirstField => then.as_mut().do_find_air_tree_node(tree_path_iter),
                     _ => panic!("Tree Path index outside tree children nodes"),
                 },
-                AirTree::DefineFunc { .. } | AirTree::DefineCyclicFuncs { .. } => unreachable!(),
+                AirTree::DefineFunc { .. }
+                | AirTree::DefineCyclicFuncs { .. }
+                | AirTree::DefineConstant { .. } => unreachable!(),
                 AirTree::FieldsEmpty {
                     constr,
                     then,