finalize aiken ir with typing

crates/aiken-lang/src/air.rs

@@ -10,6 +10,7 @@ use crate::{
 
 #[derive(Debug, Clone)]
 pub enum Air {
+    // Primitives
     Int {
         scope: Vec<u64>,
         value: String,
@@ -30,17 +31,6 @@ pub enum Air {
         value: bool,
     },
 
-    Var {
-        scope: Vec<u64>,
-        constructor: ValueConstructor,
-        name: String,
-        variant_name: String,
-    },
-
-    Fn {
-        scope: Vec<u64>,
-        params: Vec<String>,
-    },
     List {
         scope: Vec<u64>,
         count: usize,
@@ -48,48 +38,30 @@ pub enum Air {
         tail: bool,
     },
 
-    ListAccessor {
+    Tuple {
         scope: Vec<u64>,
         tipo: Arc<Type>,
-        names: Vec<String>,
+        count: usize,
-        tail: bool,
     },
 
-    ListExpose {
+    Void {
         scope: Vec<u64>,
-        tipo: Arc<Type>,
-        tail_head_names: Vec<(String, String)>,
-        tail: Option<(String, String)>,
     },
 
+    Var {
+        scope: Vec<u64>,
+        constructor: ValueConstructor,
+        name: String,
+        variant_name: String,
+    },
+
+    // Functions
     Call {
         scope: Vec<u64>,
         count: usize,
-    },
-
-    Builtin {
-        scope: Vec<u64>,
-        func: DefaultFunction,
         tipo: Arc<Type>,
     },
 
-    BinOp {
-        scope: Vec<u64>,
-        name: BinOp,
-        count: usize,
-        tipo: Arc<Type>,
-    },
-
-    Assignment {
-        scope: Vec<u64>,
-        name: String,
-    },
-
-    Assert {
-        scope: Vec<u64>,
-        constr_index: usize,
-    },
-
     DefineFunc {
         scope: Vec<u64>,
         func_name: String,
@@ -99,11 +71,42 @@ pub enum Air {
         variant_name: String,
     },
 
-    Lam {
+    Fn {
+        scope: Vec<u64>,
+        params: Vec<String>,
+    },
+
+    Builtin {
+        scope: Vec<u64>,
+        func: DefaultFunction,
+        tipo: Arc<Type>,
+    },
+
+    // Operators
+    BinOp {
+        scope: Vec<u64>,
+        name: BinOp,
+        count: usize,
+        tipo: Arc<Type>,
+    },
+
+    UnOp {
+        scope: Vec<u64>,
+        op: UnOp,
+    },
+
+    // Assignment
+    Assert {
+        scope: Vec<u64>,
+        constr_index: usize,
+    },
+
+    Let {
         scope: Vec<u64>,
         name: String,
     },
 
+    // When
     When {
         scope: Vec<u64>,
         tipo: Arc<Type>,
@@ -153,18 +156,17 @@ pub enum Air {
         inverse: bool,
     },
 
-    Discard {
-        scope: Vec<u64>,
-    },
-
     Finally {
         scope: Vec<u64>,
     },
 
+    // If
     If {
         scope: Vec<u64>,
+        tipo: Arc<Type>,
     },
 
+    // Record Creation
     Record {
         scope: Vec<u64>,
         constr_index: usize,
@@ -172,6 +174,14 @@ pub enum Air {
         count: usize,
     },
 
+    RecordUpdate {
+        scope: Vec<u64>,
+        highest_index: usize,
+        indices: Vec<(usize, Arc<Type>)>,
+        tipo: Arc<Type>,
+    },
+
+    // Field Access
     RecordAccess {
         scope: Vec<u64>,
         record_index: u64,
@@ -184,10 +194,26 @@ pub enum Air {
         indices: Vec<(usize, String, Arc<Type>)>,
     },
 
-    Tuple {
+    // ListAccess
+    ListAccessor {
         scope: Vec<u64>,
         tipo: Arc<Type>,
-        count: usize,
+        names: Vec<String>,
+        tail: bool,
+    },
+
+    ListExpose {
+        scope: Vec<u64>,
+        tipo: Arc<Type>,
+        tail_head_names: Vec<(String, String)>,
+        tail: Option<(String, String)>,
+    },
+
+    // Tuple Access
+    TupleAccessor {
+        scope: Vec<u64>,
+        names: Vec<String>,
+        tipo: Arc<Type>,
     },
 
     TupleIndex {
@@ -196,6 +222,7 @@ pub enum Air {
         tuple_index: usize,
     },
 
+    // Misc.
     Todo {
         scope: Vec<u64>,
         label: Option<String>,
@@ -213,23 +240,6 @@ pub enum Air {
         text: Option<String>,
         tipo: Arc<Type>,
     },
-
-    RecordUpdate {
-        scope: Vec<u64>,
-        highest_index: usize,
-        indices: Vec<(usize, Arc<Type>)>,
-    },
-
-    UnOp {
-        scope: Vec<u64>,
-        op: UnOp,
-    },
-
-    TupleAccessor {
-        scope: Vec<u64>,
-        names: Vec<String>,
-        tipo: Arc<Type>,
-    },
 }
 
 impl Air {
@@ -239,47 +249,93 @@ impl Air {
             | Air::String { scope, .. }
             | Air::ByteArray { scope, .. }
             | Air::Bool { scope, .. }
-            | Air::Var { scope, .. }
             | Air::List { scope, .. }
-            | Air::ListAccessor { scope, .. }
+            | Air::Tuple { scope, .. }
-            | Air::ListExpose { scope, .. }
+            | Air::Void { scope }
-            | Air::Fn { scope, .. }
+            | Air::Var { scope, .. }
             | Air::Call { scope, .. }
+            | Air::DefineFunc { scope, .. }
+            | Air::Fn { scope, .. }
             | Air::Builtin { scope, .. }
             | Air::BinOp { scope, .. }
-            | Air::Assignment { scope, .. }
+            | Air::UnOp { scope, .. }
             | Air::Assert { scope, .. }
-            | Air::DefineFunc { scope, .. }
+            | Air::Let { scope, .. }
-            | Air::Lam { scope, .. }
             | Air::When { scope, .. }
             | Air::Clause { scope, .. }
             | Air::ListClause { scope, .. }
-            | Air::TupleClause { scope, .. }
             | Air::WrapClause { scope }
+            | Air::TupleClause { scope, .. }
             | Air::ClauseGuard { scope, .. }
             | Air::ListClauseGuard { scope, .. }
-            | Air::Discard { scope }
             | Air::Finally { scope }
             | Air::If { scope, .. }
             | Air::Record { scope, .. }
+            | Air::RecordUpdate { scope, .. }
             | Air::RecordAccess { scope, .. }
             | Air::FieldsExpose { scope, .. }
-            | Air::Tuple { scope, .. }
+            | Air::ListAccessor { scope, .. }
+            | Air::ListExpose { scope, .. }
+            | Air::TupleAccessor { scope, .. }
+            | Air::TupleIndex { scope, .. }
             | Air::Todo { scope, .. }
             | Air::ErrorTerm { scope, .. }
-            | Air::RecordUpdate { scope, .. }
+            | Air::Trace { scope, .. } => scope.clone(),
-            | Air::UnOp { scope, .. }
-            | Air::Trace { scope, .. }
-            | Air::TupleAccessor { scope, .. }
-            | Air::TupleIndex { scope, .. } => scope.to_vec(),
         }
     }
 
     pub fn tipo(&self) -> Option<Arc<Type>> {
         match self {
+            Air::Int { .. } => Some(
+                Type::App {
+                    public: true,
+                    module: String::new(),
+                    name: "Int".to_string(),
+                    args: vec![],
+                }
+                .into(),
+            ),
+            Air::String { .. } => Some(
+                Type::App {
+                    public: true,
+                    module: String::new(),
+                    name: "String".to_string(),
+                    args: vec![],
+                }
+                .into(),
+            ),
+            Air::ByteArray { .. } => Some(
+                Type::App {
+                    public: true,
+                    module: String::new(),
+                    name: "ByteArray".to_string(),
+                    args: vec![],
+                }
+                .into(),
+            ),
+            Air::Bool { .. } => Some(
+                Type::App {
+                    public: true,
+                    module: String::new(),
+                    name: "Bool".to_string(),
+                    args: vec![],
+                }
+                .into(),
+            ),
+            Air::Void { .. } => Some(
+                Type::App {
+                    public: true,
+                    module: String::new(),
+                    name: "Void".to_string(),
+                    args: vec![],
+                }
+                .into(),
+            ),
+            Air::Var { constructor, .. } => Some(constructor.tipo.clone()),
+
             Air::List { tipo, .. }
-            | Air::ListAccessor { tipo, .. }
+            | Air::Tuple { tipo, .. }
-            | Air::ListExpose { tipo, .. }
+            | Air::Call { tipo, .. }
             | Air::Builtin { tipo, .. }
             | Air::BinOp { tipo, .. }
             | Air::When { tipo, .. }
@@ -287,20 +343,47 @@ impl Air {
             | Air::ListClause { tipo, .. }
             | Air::TupleClause { tipo, .. }
             | Air::ClauseGuard { tipo, .. }
+            | Air::If { tipo, .. }
             | Air::ListClauseGuard { tipo, .. }
             | Air::Record { tipo, .. }
+            | Air::RecordUpdate { tipo, .. }
             | Air::RecordAccess { tipo, .. }
-            | Air::Tuple { tipo, .. }
+            | Air::ListAccessor { tipo, .. }
+            | Air::ListExpose { tipo, .. }
+            | Air::TupleAccessor { tipo, .. }
             | Air::TupleIndex { tipo, .. }
             | Air::Todo { tipo, .. }
             | Air::ErrorTerm { tipo, .. }
-            | Air::Trace { tipo, .. }
+            | Air::Trace { tipo, .. } => Some(tipo.clone()),
-            | Air::TupleAccessor { tipo, .. }
+
-            | Air::Var {
+            Air::DefineFunc { .. }
-                constructor: ValueConstructor { tipo, .. },
+            | Air::Fn { .. }
-                ..
+            | Air::Assert { .. }
-            } => Some(tipo.clone()),
+            | Air::Let { .. }
-            _ => None,
+            | Air::WrapClause { .. }
+            | Air::Finally { .. }
+            | Air::FieldsExpose { .. } => None,
+
+            Air::UnOp { op, .. } => match op {
+                UnOp::Not => Some(
+                    Type::App {
+                        public: true,
+                        module: String::new(),
+                        name: "Bool".to_string(),
+                        args: vec![],
+                    }
+                    .into(),
+                ),
+                UnOp::Negate => Some(
+                    Type::App {
+                        public: true,
+                        module: String::new(),
+                        name: "Int".to_string(),
+                        args: vec![],
+                    }
+                    .into(),
+                ),
+            },
         }
     }
 }

crates/aiken-lang/src/builder.rs

@@ -1134,7 +1134,7 @@ pub fn monomorphize(
                 }
             }
             // TODO check on assignment if type is needed
-            Air::Assignment { .. } => {}
+            Air::Let { .. } => {}
             Air::When {
                 scope,
                 tipo,
@@ -1370,8 +1370,10 @@ pub fn monomorphize(
                 scope,
                 highest_index,
                 indices,
+                tipo,
             } => {
                 let mut new_indices = vec![];
+                let mut tipo = tipo.clone();
                 for (ind, tipo) in indices {
                     if tipo.is_generic() {
                         let mut tipo = tipo.clone();
@@ -1382,10 +1384,14 @@ pub fn monomorphize(
                         new_indices.push((ind, tipo));
                     }
                 }
+                if tipo.is_generic() {
+                    find_generics_to_replace(&mut tipo, &generic_types);
+                }
                 new_air[index] = Air::RecordUpdate {
                     scope,
                     highest_index,
                     indices: new_indices,
+                    tipo,
                 };
             }
             Air::TupleAccessor { scope, names, tipo } => {
@@ -1397,6 +1403,25 @@ pub fn monomorphize(
                     needs_variant = true;
                 }
             }
+
+            Air::Call { scope, count, tipo } => {
+                if tipo.is_generic() {
+                    let mut tipo = tipo.clone();
+                    find_generics_to_replace(&mut tipo, &generic_types);
+
+                    new_air[index] = Air::Call { scope, count, tipo };
+                    needs_variant = true;
+                }
+            }
+            Air::If { scope, tipo } => {
+                if tipo.is_generic() {
+                    let mut tipo = tipo.clone();
+                    find_generics_to_replace(&mut tipo, &generic_types);
+
+                    new_air[index] = Air::If { scope, tipo };
+                    needs_variant = true;
+                }
+            }
             _ => {}
         }
     }
@@ -1511,10 +1536,11 @@ pub fn handle_recursion_ir(
         let current_call = recursion_ir[index - 1].clone();
 
         match current_call {
-            Air::Call { scope, count } => {
+            Air::Call { scope, count, tipo } => {
                 recursion_ir[index - 1] = Air::Call {
                     scope,
                     count: count + 1,
+                    tipo,
                 }
             }
             _ => unreachable!(),

crates/aiken-lang/src/uplc.rs

@@ -242,6 +242,7 @@ impl<'a> CodeGenerator<'a> {
                 ir_stack.push(Air::Call {
                     scope: scope.clone(),
                     count: args.len(),
+                    tipo: tipo.clone(),
                 });
                 let mut scope_fun = scope.clone();
                 scope_fun.push(self.id_gen.next());
@@ -357,7 +358,7 @@ impl<'a> CodeGenerator<'a> {
                     );
 
                     if *clause_properties.needs_constr_var() {
-                        ir_stack.push(Air::Lam {
+                        ir_stack.push(Air::Let {
                             scope: scope.clone(),
                             name: constr_var.clone(),
                         });
@@ -403,6 +404,7 @@ impl<'a> CodeGenerator<'a> {
             TypedExpr::If {
                 branches,
                 final_else,
+                tipo,
                 ..
             } => {
                 let mut if_ir = vec![];
@@ -414,10 +416,12 @@ impl<'a> CodeGenerator<'a> {
                     if index == 0 {
                         if_ir.push(Air::If {
                             scope: scope.clone(),
+                            tipo: tipo.clone(),
                         });
                     } else {
                         if_ir.push(Air::If {
                             scope: branch_scope.clone(),
+                            tipo: tipo.clone(),
                         });
                     }
                     self.build_ir(&branch.condition, &mut if_ir, branch_scope.clone());
@@ -506,7 +510,9 @@ impl<'a> CodeGenerator<'a> {
                     tipo: tipo.clone(),
                 });
             }
-            TypedExpr::RecordUpdate { spread, args, .. } => {
+            TypedExpr::RecordUpdate {
+                spread, args, tipo, ..
+            } => {
                 let mut update_ir = vec![];
                 let mut spread_scope = scope.clone();
                 let mut index_types = vec![];
@@ -531,6 +537,7 @@ impl<'a> CodeGenerator<'a> {
                     scope,
                     highest_index,
                     indices: index_types,
+                    tipo: tipo.clone(),
                 });
 
                 ir_stack.append(&mut update_ir);
@@ -771,10 +778,10 @@ impl<'a> CodeGenerator<'a> {
             }
             Pattern::String { .. } => todo!(),
             Pattern::Var { name, .. } => {
-                pattern_vec.push(Air::Discard {
+                pattern_vec.push(Air::Void {
                     scope: scope.clone(),
                 });
-                pattern_vec.push(Air::Lam {
+                pattern_vec.push(Air::Let {
                     scope: scope.clone(),
                     name: name.clone(),
                 });
@@ -795,7 +802,7 @@ impl<'a> CodeGenerator<'a> {
             Pattern::VarUsage { .. } => todo!(),
             Pattern::Assign { name, pattern, .. } => {
                 let mut new_vec = vec![];
-                new_vec.push(Air::Lam {
+                new_vec.push(Air::Let {
                     scope: scope.clone(),
                     name: name.clone(),
                 });
@@ -824,7 +831,7 @@ impl<'a> CodeGenerator<'a> {
                 );
             }
             Pattern::Discard { .. } => {
-                pattern_vec.push(Air::Discard { scope });
+                pattern_vec.push(Air::Void { scope });
                 pattern_vec.append(values);
             }
             Pattern::List { elements, tail, .. } => {
@@ -837,7 +844,7 @@ impl<'a> CodeGenerator<'a> {
                 }
                 *clause_properties.needs_constr_var() = false;
 
-                pattern_vec.push(Air::Discard {
+                pattern_vec.push(Air::Void {
                     scope: scope.clone(),
                 });
 
@@ -965,7 +972,7 @@ impl<'a> CodeGenerator<'a> {
             Pattern::VarUsage { .. } => todo!(),
             Pattern::Assign { .. } => todo!(),
             Pattern::Discard { .. } => {
-                pattern_vec.push(Air::Discard { scope });
+                pattern_vec.push(Air::Void { scope });
 
                 pattern_vec.append(values);
             }
@@ -1198,7 +1205,7 @@ impl<'a> CodeGenerator<'a> {
 
                     let pattern_type = &tipo.get_inner_types()[index];
 
-                    pattern_vec.push(Air::Lam {
+                    pattern_vec.push(Air::Let {
                         scope: scope.clone(),
                         name: new_name.clone(),
                     });
@@ -1254,7 +1261,7 @@ impl<'a> CodeGenerator<'a> {
                         inverse: false,
                     });
 
-                    pattern_vec.push(Air::Discard { scope });
+                    pattern_vec.push(Air::Void { scope });
                 } else {
                     for (index, _) in elements.iter().enumerate() {
                         let prev_tail_name = if index == 0 {
@@ -1300,7 +1307,7 @@ impl<'a> CodeGenerator<'a> {
                                     inverse: true,
                                 });
 
-                                pattern_vec.push(Air::Discard {
+                                pattern_vec.push(Air::Void {
                                     scope: scope.clone(),
                                 });
 
@@ -1445,7 +1452,7 @@ impl<'a> CodeGenerator<'a> {
         match pattern {
             Pattern::Int { .. } | Pattern::String { .. } => unreachable!(),
             Pattern::Var { name, .. } => {
-                pattern_vec.push(Air::Assignment {
+                pattern_vec.push(Air::Let {
                     name: name.clone(),
                     scope,
                 });
@@ -1455,7 +1462,7 @@ impl<'a> CodeGenerator<'a> {
             Pattern::VarUsage { .. } => todo!(),
             Pattern::Assign { .. } => todo!(),
             Pattern::Discard { .. } => {
-                pattern_vec.push(Air::Assignment {
+                pattern_vec.push(Air::Let {
                     name: "_".to_string(),
                     scope,
                 });
@@ -1485,7 +1492,7 @@ impl<'a> CodeGenerator<'a> {
                     }
                     AssignmentKind::Assert => {
                         let name_id = self.id_gen.next();
-                        pattern_vec.push(Air::Lam {
+                        pattern_vec.push(Air::Let {
                             scope: scope.clone(),
                             name: format!("__constr_{}", name_id),
                         });
@@ -1551,7 +1558,7 @@ impl<'a> CodeGenerator<'a> {
             Pattern::VarUsage { .. } => todo!(),
             Pattern::Assign { .. } => todo!(),
             Pattern::Discard { .. } => {
-                pattern_vec.push(Air::Discard { scope });
+                pattern_vec.push(Air::Void { scope });
 
                 pattern_vec.append(values);
             }
@@ -2562,6 +2569,7 @@ impl<'a> CodeGenerator<'a> {
                     highest_index,
                     indices,
                     scope,
+                    tipo,
                 } => {
                     let mut new_indices = vec![];
                     for (ind, tipo) in indices {
@@ -2574,6 +2582,7 @@ impl<'a> CodeGenerator<'a> {
                         scope,
                         indices: new_indices,
                         highest_index,
+                        tipo,
                     };
                 }
                 Air::Record {
@@ -2636,7 +2645,7 @@ impl<'a> CodeGenerator<'a> {
                     let record_type = record.tipo();
                     if let Some(record_type) = record_type {
                         if check_replaceable_opaque_type(&record_type, self.data_types) {
-                            ir_stack[index] = Air::Lam {
+                            ir_stack[index] = Air::Let {
                                 scope,
                                 name: indices[0].1.clone(),
                             };
@@ -2669,6 +2678,25 @@ impl<'a> CodeGenerator<'a> {
                         };
                     }
                 }
+                Air::Call { scope, count, tipo } => {
+                    let mut replaced_type = tipo.clone();
+                    replace_opaque_type(&mut replaced_type, self.data_types.clone());
+
+                    ir_stack[index] = Air::Call {
+                        scope,
+                        tipo: replaced_type,
+                        count,
+                    };
+                }
+                Air::If { scope, tipo } => {
+                    let mut replaced_type = tipo.clone();
+                    replace_opaque_type(&mut replaced_type, self.data_types.clone());
+
+                    ir_stack[index] = Air::If {
+                        scope,
+                        tipo: replaced_type,
+                    };
+                }
                 _ => {}
             }
         }
@@ -2868,7 +2896,7 @@ impl<'a> CodeGenerator<'a> {
                     }
                 };
             }
-            Air::Discard { .. } => {
+            Air::Void { .. } => {
                 arg_stack.push(Term::Constant(UplcConstant::Unit));
             }
             Air::List {
@@ -3461,23 +3489,6 @@ impl<'a> CodeGenerator<'a> {
                 };
                 arg_stack.push(term);
             }
-            Air::Assignment { name, .. } => {
-                let right_hand = arg_stack.pop().unwrap();
-                let lam_body = arg_stack.pop().unwrap();
-
-                let term = apply_wrap(
-                    Term::Lambda {
-                        parameter_name: Name {
-                            text: name,
-                            unique: 0.into(),
-                        },
-                        body: lam_body.into(),
-                    },
-                    right_hand,
-                );
-
-                arg_stack.push(term);
-            }
             Air::Assert { constr_index, .. } => {
                 let constr = arg_stack.pop().unwrap();
 
@@ -3586,7 +3597,7 @@ impl<'a> CodeGenerator<'a> {
                     arg_stack.push(term);
                 }
             }
-            Air::Lam { name, .. } => {
+            Air::Let { name, .. } => {
                 let arg = arg_stack.pop().unwrap();
 
                 let mut term = arg_stack.pop().unwrap();