feat: implement most of airtree build

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

@@ -153,6 +153,7 @@ pub enum AirTree {
     // If
     If {
         tipo: Arc<Type>,
+        pattern: Box<AirTree>,
         then: Box<AirTree>,
         otherwise: Box<AirTree>,
     },
@@ -216,6 +217,9 @@ pub enum AirTree {
         then: Box<AirTree>,
     },
     NoOp,
+    Sequence {
+        expressions: Vec<AirTree>,
+    },
     FieldsEmpty {
         constr: Box<AirTree>,
     },
@@ -241,8 +245,22 @@ impl AirTree {
     pub fn bool(value: bool) -> AirTree {
         AirTree::Bool { value }
     }
-    pub fn list(items: Vec<AirTree>, tipo: Arc<Type>, tail: bool) -> AirTree {
+    pub fn list(mut items: Vec<AirTree>, tipo: Arc<Type>, tail: Option<AirTree>) -> AirTree {
-        AirTree::List { tipo, tail, items }
+        if let Some(tail) = tail {
+            items.push(tail);
+
+            AirTree::List {
+                tipo,
+                tail: true,
+                items,
+            }
+        } else {
+            AirTree::List {
+                tipo,
+                tail: false,
+                items,
+            }
+        }
     }
     pub fn tuple(items: Vec<AirTree>, tipo: Arc<Type>) -> AirTree {
         AirTree::Tuple { tipo, items }
@@ -450,20 +468,26 @@ impl AirTree {
             then: then.into(),
         }
     }
-    pub fn if_branches(mut branches: Vec<AirTree>, tipo: Arc<Type>, otherwise: AirTree) -> AirTree {
+    pub fn if_branches(
+        mut branches: Vec<(AirTree, AirTree)>,
+        tipo: Arc<Type>,
+        otherwise: AirTree,
+    ) -> AirTree {
         assert!(branches.len() > 0);
         let last_if = branches.pop().unwrap();
 
         let mut final_if = AirTree::If {
             tipo: tipo.clone(),
-            then: last_if.into(),
+            pattern: last_if.0.into(),
+            then: last_if.1.into(),
             otherwise: otherwise.into(),
         };
 
         while let Some(branch) = branches.pop() {
             final_if = AirTree::If {
                 tipo: tipo.clone(),
-                then: branch.into(),
+                pattern: branch.0.into(),
+                then: branch.1.into(),
                 otherwise: final_if.into(),
             };
         }
@@ -569,6 +593,9 @@ 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(),
@@ -672,6 +699,7 @@ impl AirTree {
             AirTree::Finally { pattern, then } => todo!(),
             AirTree::If {
                 tipo,
+                pattern,
                 then,
                 otherwise,
             } => todo!(),
@@ -720,6 +748,7 @@ 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

@@ -0,0 +1,321 @@
+mod builder;
+
+use std::{rc::Rc, sync::Arc};
+
+use indexmap::{IndexMap, IndexSet};
+use itertools::Itertools;
+use uplc::ast::{Name, Program, Term};
+
+use crate::{
+    ast::{TypedDataType, TypedFunction, TypedValidator},
+    expr::TypedExpr,
+    gen_uplc::{
+        air::Air,
+        builder::{self as build, DataTypeKey, FunctionAccessKey},
+        tree::AirTree,
+        CodeGenFunction,
+    },
+    tipo::{ModuleValueConstructor, TypeInfo, ValueConstructor, ValueConstructorVariant},
+};
+
+#[derive(Clone)]
+pub struct CodeGenerator<'a> {
+    defined_functions: IndexMap<FunctionAccessKey, ()>,
+    functions: IndexMap<FunctionAccessKey, &'a TypedFunction>,
+    data_types: IndexMap<DataTypeKey, &'a TypedDataType>,
+    module_types: IndexMap<&'a String, &'a TypeInfo>,
+    needs_field_access: bool,
+    code_gen_functions: IndexMap<String, CodeGenFunction>,
+    zero_arg_functions: IndexMap<FunctionAccessKey, Vec<Air>>,
+    tracing: bool,
+}
+
+impl<'a> CodeGenerator<'a> {
+    pub fn new(
+        functions: IndexMap<FunctionAccessKey, &'a TypedFunction>,
+        data_types: IndexMap<DataTypeKey, &'a TypedDataType>,
+        module_types: IndexMap<&'a String, &'a TypeInfo>,
+        tracing: bool,
+    ) -> Self {
+        CodeGenerator {
+            defined_functions: IndexMap::new(),
+            functions,
+            data_types,
+            module_types,
+            needs_field_access: false,
+            code_gen_functions: IndexMap::new(),
+            zero_arg_functions: IndexMap::new(),
+            tracing,
+        }
+    }
+
+    pub fn reset(&mut self) {
+        self.code_gen_functions = IndexMap::new();
+        self.zero_arg_functions = IndexMap::new();
+        self.needs_field_access = false;
+        self.defined_functions = IndexMap::new();
+    }
+
+    pub fn generate(
+        &mut self,
+        TypedValidator {
+            fun,
+            other_fun,
+            params,
+            ..
+        }: &TypedValidator,
+    ) -> Program<Name> {
+        todo!()
+    }
+
+    pub fn generate_test(&mut self, test_body: &TypedExpr) -> Program<Name> {
+        todo!()
+    }
+
+    fn finalize(&mut self, term: Term<Name>) -> Program<Name> {
+        todo!()
+    }
+
+    fn build(&mut self, body: &TypedExpr) -> AirTree {
+        match body {
+            TypedExpr::Int { value, .. } => AirTree::int(value),
+            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(),
+                )
+            }
+
+            TypedExpr::Var {
+                constructor, name, ..
+            } => AirTree::var(constructor.clone(), name, ""),
+
+            TypedExpr::Fn { args, body, .. } => AirTree::anon_func(
+                args.iter()
+                    .map(|arg| arg.arg_name.get_variable_name().unwrap_or("_").to_string())
+                    .collect_vec(),
+                self.build(body),
+            ),
+
+            TypedExpr::List {
+                tipo,
+                elements,
+                tail,
+                ..
+            } => AirTree::list(
+                elements.iter().map(|elem| self.build(elem)).collect_vec(),
+                tipo.clone(),
+                tail.as_ref().map(|tail| self.build(tail)),
+            ),
+
+            TypedExpr::Call {
+                tipo, fun, args, ..
+            } => match fun.as_ref() {
+                TypedExpr::Var {
+                    constructor:
+                        ValueConstructor {
+                            variant:
+                                ValueConstructorVariant::Record {
+                                    name: constr_name, ..
+                                },
+                            ..
+                        },
+                    ..
+                }
+                | TypedExpr::ModuleSelect {
+                    constructor:
+                        ModuleValueConstructor::Record {
+                            name: constr_name, ..
+                        },
+                    ..
+                } => {
+                    let Some(data_type) = build::lookup_data_type_by_tipo(&self.data_types, tipo)
+                    else {unreachable!("Creating a record with no record definition.")};
+
+                    let (constr_index, _) = data_type
+                        .constructors
+                        .iter()
+                        .enumerate()
+                        .find(|(_, dt)| &dt.name == constr_name)
+                        .unwrap();
+
+                    let constr_args = args.iter().map(|arg| self.build(&arg.value)).collect_vec();
+
+                    AirTree::create_constr(constr_index, tipo.clone(), constr_args)
+                }
+
+                TypedExpr::Var {
+                    constructor:
+                        ValueConstructor {
+                            variant: ValueConstructorVariant::ModuleFn { builtin, .. },
+                            ..
+                        },
+                    ..
+                } => {
+                    let Some(fun_arg_types) = fun.tipo().arg_types() else {unreachable!("Expected a function type with arguments")};
+
+                    let func_args = args
+                        .iter()
+                        .zip(fun_arg_types)
+                        .map(|(arg, arg_tipo)| {
+                            let mut arg_val = self.build(&arg.value);
+
+                            if arg_tipo.is_data() && !arg.value.tipo().is_data() {
+                                arg_val = AirTree::wrap_data(arg_val, arg.value.tipo())
+                            }
+                            arg_val
+                        })
+                        .collect_vec();
+
+                    if let Some(func) = builtin {
+                        AirTree::builtin(*func, tipo.clone(), func_args)
+                    } else {
+                        AirTree::call(self.build(fun.as_ref()), tipo.clone(), func_args)
+                    }
+                }
+
+                TypedExpr::ModuleSelect {
+                    tipo,
+                    module_name,
+                    constructor: ModuleValueConstructor::Fn { name, .. },
+                    ..
+                } => {
+                    let type_info = self.module_types.get(module_name).unwrap();
+                    let value = type_info.values.get(name).unwrap();
+
+                    let ValueConstructorVariant::ModuleFn { builtin, .. } = &value.variant else {unreachable!("Missing module function definition")};
+
+                    let Some(fun_arg_types) = fun.tipo().arg_types() else {unreachable!("Expected a function type with arguments")};
+
+                    let func_args = args
+                        .iter()
+                        .zip(fun_arg_types)
+                        .map(|(arg, arg_tipo)| {
+                            let mut arg_val = self.build(&arg.value);
+
+                            if arg_tipo.is_data() && !arg.value.tipo().is_data() {
+                                arg_val = AirTree::wrap_data(arg_val, arg.value.tipo())
+                            }
+                            arg_val
+                        })
+                        .collect_vec();
+
+                    if let Some(func) = builtin {
+                        AirTree::builtin(*func, tipo.clone(), func_args)
+                    } else {
+                        AirTree::call(self.build(fun.as_ref()), tipo.clone(), func_args)
+                    }
+                }
+                _ => todo!("IS THIS REACHABLE?"),
+            },
+            TypedExpr::BinOp {
+                name,
+                tipo,
+                left,
+                right,
+                ..
+            } => AirTree::binop(*name, tipo.clone(), self.build(left), self.build(right)),
+
+            TypedExpr::Assignment {
+                tipo,
+                value,
+                pattern,
+                kind,
+                ..
+            } => {
+                let mut replaced_type = tipo.clone();
+                build::replace_opaque_type(&mut replaced_type, &self.data_types);
+
+                let air_value = self.build(value);
+
+                builder::assignment_air_tree(
+                    pattern,
+                    air_value,
+                    tipo,
+                    build::AssignmentProperties {
+                        value_type: value.tipo(),
+                        kind: *kind,
+                    },
+                )
+            }
+
+            TypedExpr::Trace {
+                tipo, then, text, ..
+            } => AirTree::trace(self.build(text), tipo.clone(), self.build(then)),
+
+            TypedExpr::When { .. } => todo!(),
+
+            TypedExpr::If {
+                branches,
+                final_else,
+                tipo,
+                ..
+            } => AirTree::if_branches(
+                branches
+                    .iter()
+                    .map(|branch| (self.build(&branch.condition), self.build(&branch.body)))
+                    .collect_vec(),
+                tipo.clone(),
+                self.build(final_else),
+            ),
+
+            TypedExpr::RecordAccess {
+                tipo,
+                index,
+                record,
+                ..
+            } => AirTree::record_access(*index, tipo.clone(), self.build(record)),
+
+            TypedExpr::ModuleSelect { .. } => todo!(),
+
+            TypedExpr::Tuple { tipo, elems, .. } => AirTree::tuple(
+                elems.iter().map(|elem| self.build(elem)).collect_vec(),
+                tipo.clone(),
+            ),
+
+            TypedExpr::TupleIndex {
+                tipo, index, tuple, ..
+            } => AirTree::tuple_index(*index, tipo.clone(), self.build(tuple)),
+
+            TypedExpr::ErrorTerm { tipo, .. } => AirTree::error(tipo.clone()),
+
+            TypedExpr::RecordUpdate {
+                tipo, spread, args, ..
+            } => {
+                let mut index_types = vec![];
+                let mut update_args = vec![];
+
+                let mut highest_index = 0;
+                let record = self.build(spread);
+
+                for arg in args
+                    .iter()
+                    .sorted_by(|arg1, arg2| arg1.index.cmp(&arg2.index))
+                {
+                    let arg_val = self.build(&arg.value);
+
+                    if arg.index > highest_index {
+                        highest_index = arg.index;
+                    }
+
+                    index_types.push((arg.index, arg.value.tipo()));
+                    update_args.push(arg_val);
+                }
+
+                AirTree::record_update(
+                    index_types,
+                    highest_index,
+                    tipo.clone(),
+                    record,
+                    update_args,
+                )
+            }
+
+            TypedExpr::UnOp { value, op, .. } => AirTree::unop(*op, self.build(value)),
+        }
+    }
+}

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

@@ -0,0 +1,19 @@
+use std::sync::Arc;
+
+use crate::{
+    ast::Pattern,
+    gen_uplc::{builder::AssignmentProperties, tree::AirTree},
+    tipo::{PatternConstructor, Type},
+};
+
+pub fn assignment_air_tree(
+    pattern: &Pattern<PatternConstructor, Arc<Type>>,
+    value: AirTree,
+    tipo: &Arc<Type>,
+    props: AssignmentProperties,
+) -> AirTree {
+    todo!()
+}
+
+
+

crates/aiken-lang/src/lib.rs

@@ -8,6 +8,7 @@ pub mod builtins;
 pub mod expr;
 pub mod format;
 pub mod gen_uplc;
+pub mod gen_uplc2;
 pub mod levenshtein;
 pub mod parser;
 pub mod pretty;