feat: update tree to allow for let hoisting

feat: start on build for when expressions
feat: add builder methods for AirTree

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

@@ -78,7 +78,7 @@ pub enum AirTree {
     Let {
         name: String,
         value: Box<AirTree>,
-        hoisted_over: Box<AirTree>,
+        hoisted_over: Box<Option<AirTree>>,
     },
     UnWrapData {
         tipo: Arc<Type>,
@@ -217,9 +217,6 @@ pub enum AirTree {
         then: Box<AirTree>,
     },
     NoOp,
-    Sequence {
-        expressions: Vec<AirTree>,
-    },
     FieldsEmpty {
         constr: Box<AirTree>,
     },
@@ -328,11 +325,11 @@ impl AirTree {
             arg: arg.into(),
         }
     }
-    pub fn let_assignment(name: impl ToString, value: AirTree, hoisting_over: AirTree) -> AirTree {
+    pub fn let_assignment(name: impl ToString, value: AirTree) -> AirTree {
         AirTree::Let {
             name: name.to_string(),
             value: value.into(),
-            hoisted_over: hoisting_over.into(),
+            hoisted_over: None.into(),
         }
     }
     pub fn unwrap_data(value: AirTree, tipo: Arc<Type>) -> AirTree {
@@ -473,7 +470,7 @@ impl AirTree {
         tipo: Arc<Type>,
         otherwise: AirTree,
     ) -> AirTree {
-        assert!(branches.len() > 0);
+        assert!(!branches.is_empty());
         let last_if = branches.pop().unwrap();
 
         let mut final_if = AirTree::If {
@@ -593,9 +590,6 @@ impl AirTree {
     pub fn no_op() -> AirTree {
         AirTree::NoOp
     }
-    pub fn sequence(expressions: Vec<AirTree>) -> AirTree {
-        AirTree::Sequence { expressions }
-    }
     pub fn fields_empty(constr: AirTree) -> AirTree {
         AirTree::FieldsEmpty {
             constr: constr.into(),
@@ -604,6 +598,16 @@ impl AirTree {
     pub fn list_empty(list: AirTree) -> AirTree {
         AirTree::ListEmpty { list: list.into() }
     }
+    pub fn hoist_let(assignment: AirTree, next_exp: AirTree) -> AirTree {
+        match assignment {
+            AirTree::Let { name, value, .. } => AirTree::Let {
+                name,
+                value,
+                hoisted_over: Some(next_exp).into(),
+            },
+            _ => unimplemented!("IS THIS REACHABLE?"),
+        }
+    }
 
     pub fn to_air_vec(tree: AirTree) -> Vec<Air> {
         match tree {
@@ -748,7 +752,6 @@ impl AirTree {
             AirTree::ErrorTerm { tipo } => todo!(),
             AirTree::Trace { tipo, msg, then } => todo!(),
             AirTree::NoOp => todo!(),
-            AirTree::Sequence { .. } => todo!(),
             AirTree::FieldsEmpty { constr } => todo!(),
             AirTree::ListEmpty { list } => todo!(),
         }

crates/aiken-lang/src/gen_uplc2.rs

@@ -7,11 +7,11 @@ use itertools::Itertools;
 use uplc::ast::{Name, Program, Term};
 
 use crate::{
-    ast::{TypedDataType, TypedFunction, TypedValidator},
+    ast::{AssignmentKind, Span, TypedDataType, TypedFunction, TypedValidator},
     expr::TypedExpr,
     gen_uplc::{
         air::Air,
-        builder::{self as build, DataTypeKey, FunctionAccessKey},
+        builder::{self as build, AssignmentProperties, DataTypeKey, FunctionAccessKey},
         tree::AirTree,
         CodeGenFunction,
     },
@@ -82,12 +82,16 @@ impl<'a> CodeGenerator<'a> {
             TypedExpr::String { value, .. } => AirTree::string(value),
             TypedExpr::ByteArray { bytes, .. } => AirTree::byte_array(bytes.clone()),
             TypedExpr::Sequence { expressions, .. } | TypedExpr::Pipeline { expressions, .. } => {
-                AirTree::sequence(
-                    expressions
-                        .iter()
-                        .map(|expression| self.build(expression))
-                        .collect_vec(),
-                )
+                let mut expressions = expressions.clone();
+
+                let mut last_exp = self.build(&expressions.pop().unwrap());
+
+                while let Some(expression) = expressions.pop() {
+                    let exp_tree = self.build(&expression);
+
+                    last_exp = AirTree::hoist_let(exp_tree, last_exp);
+                }
+                last_exp
             }
 
             TypedExpr::Var {
@@ -236,7 +240,7 @@ impl<'a> CodeGenerator<'a> {
                     pattern,
                     air_value,
                     tipo,
-                    build::AssignmentProperties {
+                    AssignmentProperties {
                         value_type: value.tipo(),
                         kind: *kind,
                     },
@@ -247,7 +251,46 @@ impl<'a> CodeGenerator<'a> {
                 tipo, then, text, ..
             } => AirTree::trace(self.build(text), tipo.clone(), self.build(then)),
 
-            TypedExpr::When { .. } => todo!(),
+            TypedExpr::When {
+                tipo,
+                subject,
+                clauses,
+                ..
+            } => {
+                let mut clauses = clauses.clone();
+
+                if clauses.is_empty() {
+                    panic!("We should have one clause at least")
+                } else if clauses.len() == 1 {
+                    let last_clause = clauses.pop().unwrap();
+
+                    let clause_then = self.build(&last_clause.then);
+
+                    let subject_val = self.build(subject);
+
+                    let assignment = builder::assignment_air_tree(
+                        &last_clause.pattern,
+                        subject_val,
+                        tipo,
+                        AssignmentProperties {
+                            value_type: subject.tipo(),
+                            kind: AssignmentKind::Let,
+                        },
+                    );
+
+                    AirTree::hoist_let(assignment, clause_then)
+                } else {
+                    clauses = if subject.tipo().is_list() {
+                        build::rearrange_clauses(clauses.clone())
+                    } else {
+                        clauses
+                    };
+
+                    let last_clause = clauses.pop().unwrap();
+
+                    todo!()
+                }
+            }
 
             TypedExpr::If {
                 branches,
@@ -270,7 +313,63 @@ impl<'a> CodeGenerator<'a> {
                 ..
             } => AirTree::record_access(*index, tipo.clone(), self.build(record)),
 
-            TypedExpr::ModuleSelect { .. } => todo!(),
+            TypedExpr::ModuleSelect {
+                tipo,
+                module_name,
+                constructor,
+                ..
+            } => match constructor {
+                ModuleValueConstructor::Record {
+                    name,
+                    arity,
+                    tipo,
+                    field_map,
+                    ..
+                } => {
+                    let data_type = build::lookup_data_type_by_tipo(&self.data_types, tipo);
+
+                    let val_constructor = ValueConstructor::public(
+                        tipo.clone(),
+                        ValueConstructorVariant::Record {
+                            name: name.clone(),
+                            arity: *arity,
+                            field_map: field_map.clone(),
+                            location: Span::empty(),
+                            module: module_name.clone(),
+                            constructors_count: data_type.unwrap().constructors.len() as u16,
+                        },
+                    );
+
+                    AirTree::var(val_constructor, name, "")
+                }
+                ModuleValueConstructor::Fn { name, module, .. } => {
+                    let func = self.functions.get(&FunctionAccessKey {
+                        module_name: module_name.clone(),
+                        function_name: name.clone(),
+                        variant_name: String::new(),
+                    });
+
+                    let type_info = self.module_types.get(module_name).unwrap();
+                    let value = type_info.values.get(name).unwrap();
+
+                    if let Some(_func) = func {
+                        AirTree::var(
+                            ValueConstructor::public(tipo.clone(), value.variant.clone()),
+                            format!("{module}_{name}"),
+                            "",
+                        )
+                    } else {
+                        let ValueConstructorVariant::ModuleFn {
+                            builtin: Some(builtin), ..
+                        } = &value.variant else {
+                            unreachable!("Didn't find the function definition.")
+                        };
+
+                        AirTree::builtin(*builtin, tipo.clone(), vec![])
+                    }
+                }
+                ModuleValueConstructor::Constant { literal, .. } => builder::constants_ir(literal),
+            },
 
             TypedExpr::Tuple { tipo, elems, .. } => AirTree::tuple(
                 elems.iter().map(|elem| self.build(elem)).collect_vec(),

crates/aiken-lang/src/gen_uplc2/builder.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use crate::{
-    ast::Pattern,
+    ast::{Constant, Pattern},
     gen_uplc::{builder::AssignmentProperties, tree::AirTree},
     tipo::{PatternConstructor, Type},
 };
@@ -15,5 +15,21 @@ pub fn assignment_air_tree(
     todo!()
 }
 
+pub fn handle_each_clause(
+    pattern: &Pattern<PatternConstructor, Arc<Type>>,
+    value: AirTree,
+    tipo: &Arc<Type>,
+    props: AssignmentProperties,
+) -> AirTree {
+    todo!()
+}
+
+pub fn constants_ir(literal: &Constant) -> AirTree {
+    match literal {
+        Constant::Int { value, .. } => AirTree::int(value),
+        Constant::String { value, .. } => AirTree::string(value),
+        Constant::ByteArray { bytes, .. } => AirTree::byte_array(bytes.clone()),
+    }
+}