feat: Add codegen for handling clause guards

crates/aiken-lang/src/builder.rs

@@ -18,9 +18,12 @@ use uplc::{
 
 use crate::{
     air::Air,
-    ast::{Clause, Constant, DataType, Pattern, Span, TypedArg, TypedDataType},
+    ast::{
+        BinOp, Clause, ClauseGuard, Constant, DataType, Pattern, Span, TypedArg, TypedDataType,
+        UnOp,
+    },
     expr::TypedExpr,
-    tipo::{PatternConstructor, Type, TypeVar, ValueConstructorVariant},
+    tipo::{PatternConstructor, Type, TypeVar, ValueConstructor, ValueConstructorVariant},
 };
 
 #[derive(Clone, Debug)]
@@ -1625,3 +1628,116 @@ pub fn replace_opaque_type(t: &mut Arc<Type>, data_types: HashMap<DataTypeKey, &
         }
     }
 }
+
+pub fn handle_clause_guard(
+    clause_guard: &ClauseGuard<Arc<Type>, String>,
+    clause_guard_vec: &mut Vec<Air>,
+    scope: Vec<u64>,
+) {
+    match clause_guard {
+        ClauseGuard::Not { value, .. } => {
+            clause_guard_vec.push(Air::UnOp {
+                scope: scope.clone(),
+                op: UnOp::Not,
+            });
+
+            handle_clause_guard(value, clause_guard_vec, scope);
+        }
+        ClauseGuard::Equals { left, right, .. } => {
+            clause_guard_vec.push(Air::BinOp {
+                scope: scope.clone(),
+                name: BinOp::Eq,
+                count: 2,
+                tipo: left.tipo(),
+            });
+            handle_clause_guard(left, clause_guard_vec, scope.clone());
+            handle_clause_guard(right, clause_guard_vec, scope);
+        }
+        ClauseGuard::NotEquals { left, right, .. } => {
+            clause_guard_vec.push(Air::BinOp {
+                scope: scope.clone(),
+                name: BinOp::NotEq,
+                count: 2,
+                tipo: left.tipo(),
+            });
+            handle_clause_guard(left, clause_guard_vec, scope.clone());
+            handle_clause_guard(right, clause_guard_vec, scope);
+        }
+        ClauseGuard::GtInt { left, right, .. } => {
+            clause_guard_vec.push(Air::BinOp {
+                scope: scope.clone(),
+                name: BinOp::GtInt,
+                count: 2,
+                tipo: left.tipo(),
+            });
+            handle_clause_guard(left, clause_guard_vec, scope.clone());
+            handle_clause_guard(right, clause_guard_vec, scope);
+        }
+        ClauseGuard::GtEqInt { left, right, .. } => {
+            clause_guard_vec.push(Air::BinOp {
+                scope: scope.clone(),
+                name: BinOp::GtEqInt,
+                count: 2,
+                tipo: left.tipo(),
+            });
+            handle_clause_guard(left, clause_guard_vec, scope.clone());
+            handle_clause_guard(right, clause_guard_vec, scope);
+        }
+        ClauseGuard::LtInt { left, right, .. } => {
+            clause_guard_vec.push(Air::BinOp {
+                scope: scope.clone(),
+                name: BinOp::LtInt,
+                count: 2,
+                tipo: left.tipo(),
+            });
+            handle_clause_guard(left, clause_guard_vec, scope.clone());
+            handle_clause_guard(right, clause_guard_vec, scope);
+        }
+        ClauseGuard::LtEqInt { left, right, .. } => {
+            clause_guard_vec.push(Air::BinOp {
+                scope: scope.clone(),
+                name: BinOp::LtEqInt,
+                count: 2,
+                tipo: left.tipo(),
+            });
+            handle_clause_guard(left, clause_guard_vec, scope.clone());
+            handle_clause_guard(right, clause_guard_vec, scope);
+        }
+        ClauseGuard::Or { left, right, .. } => {
+            clause_guard_vec.push(Air::BinOp {
+                scope: scope.clone(),
+                name: BinOp::Or,
+                count: 2,
+                tipo: left.tipo(),
+            });
+            handle_clause_guard(left, clause_guard_vec, scope.clone());
+            handle_clause_guard(right, clause_guard_vec, scope);
+        }
+        ClauseGuard::And { left, right, .. } => {
+            clause_guard_vec.push(Air::BinOp {
+                scope: scope.clone(),
+                name: BinOp::And,
+                count: 2,
+                tipo: left.tipo(),
+            });
+            handle_clause_guard(left, clause_guard_vec, scope.clone());
+            handle_clause_guard(right, clause_guard_vec, scope);
+        }
+        ClauseGuard::Var { tipo, name, .. } => {
+            clause_guard_vec.push(Air::Var {
+                scope,
+                constructor: ValueConstructor::public(
+                    tipo.clone(),
+                    ValueConstructorVariant::LocalVariable {
+                        location: Span::empty(),
+                    },
+                ),
+                name: name.clone(),
+                variant_name: String::new(),
+            });
+        }
+        ClauseGuard::Constant(constant) => {
+            constants_ir(constant, clause_guard_vec, scope);
+        }
+    }
+}

crates/aiken-lang/src/uplc.rs

@@ -28,10 +28,12 @@ use crate::{
     builder::{
         check_replaceable_opaque_type, check_when_pattern_needs, constants_ir,
         convert_constants_to_data, convert_data_to_type, convert_type_to_data, get_common_ancestor,
-        get_generics_and_type, handle_func_deps_ir, handle_recursion_ir, list_access_to_uplc,
+        get_generics_and_type, handle_clause_guard, handle_func_deps_ir, handle_recursion_ir,
-        lookup_data_type_by_tipo, monomorphize, rearrange_clauses, replace_opaque_type,
+        list_access_to_uplc, lookup_data_type_by_tipo, monomorphize, rearrange_clauses,
-        wrap_validator_args, ClauseProperties, DataTypeKey, FuncComponents, FunctionAccessKey,
+        replace_opaque_type, wrap_validator_args, ClauseProperties, DataTypeKey, FuncComponents,
+        FunctionAccessKey,
     },
+    builtins::bool,
     expr::TypedExpr,
     tipo::{
         self, ModuleValueConstructor, PatternConstructor, Type, TypeInfo, ValueConstructor,
@@ -621,6 +623,33 @@ impl<'a> CodeGenerator<'a> {
 
             self.build_ir(&clause.then, &mut clause_then_vec, scope.clone());
 
+            if let Some(clause_guard) = &clause.guard {
+                let mut clause_guard_vec = vec![];
+                *clause_properties.is_complex_clause() = true;
+                let clause_guard_name = format!("__clause_guard_{}", self.id_gen.next());
+
+                clause_guard_vec.push(Air::Lam {
+                    scope: scope.clone(),
+                    name: clause_guard_name.clone(),
+                });
+
+                handle_clause_guard(clause_guard, &mut clause_guard_vec, scope.clone());
+
+                clause_guard_vec.push(Air::ClauseGuard {
+                    scope: scope.clone(),
+                    subject_name: clause_guard_name,
+                    tipo: bool(),
+                });
+
+                clause_guard_vec.push(Air::Bool {
+                    scope: scope.clone(),
+                    value: true,
+                });
+
+                clause_guard_vec.append(&mut clause_then_vec);
+                clause_then_vec = clause_guard_vec;
+            }
+
             match clause_properties {
                 ClauseProperties::ConstrClause {
                     original_subject_name,
@@ -3644,7 +3673,47 @@ impl<'a> CodeGenerator<'a> {
                 let mut term = arg_stack.pop().unwrap();
 
                 if tipo.is_bool() {
+                    if complex_clause {
+                        let other_clauses = term;
                         if matches!(clause, Term::Constant(UplcConstant::Bool(true))) {
+                            term = if_else(
+                                Term::Var(Name {
+                                    text: subject_name,
+                                    unique: 0.into(),
+                                }),
+                                Term::Delay(body.into()),
+                                Term::Var(Name {
+                                    text: "__other_clauses_delayed".to_string(),
+                                    unique: 0.into(),
+                                }),
+                            )
+                            .force_wrap();
+                        } else {
+                            term = if_else(
+                                Term::Var(Name {
+                                    text: subject_name,
+                                    unique: 0.into(),
+                                }),
+                                Term::Var(Name {
+                                    text: "__other_clauses_delayed".to_string(),
+                                    unique: 0.into(),
+                                }),
+                                Term::Delay(body.into()),
+                            )
+                            .force_wrap();
+                        }
+
+                        term = apply_wrap(
+                            Term::Lambda {
+                                parameter_name: Name {
+                                    text: "__other_clauses_delayed".to_string(),
+                                    unique: 0.into(),
+                                },
+                                body: term.into(),
+                            },
+                            Term::Delay(other_clauses.into()),
+                        );
+                    } else if matches!(clause, Term::Constant(UplcConstant::Bool(true))) {
                         term = delayed_if_else(
                             Term::Var(Name {
                                 text: subject_name,
@@ -4635,6 +4704,21 @@ impl<'a> CodeGenerator<'a> {
 
                 let tuple_types = tipo.get_inner_types();
 
+                if complex_clause {
+                    let next_clause = arg_stack.pop().unwrap();
+
+                    term = apply_wrap(
+                        Term::Lambda {
+                            parameter_name: Name {
+                                text: "__other_clauses_delayed".to_string(),
+                                unique: 0.into(),
+                            },
+                            body: term.into(),
+                        },
+                        Term::Delay(next_clause.into()),
+                    )
+                }
+
                 if tuple_types.len() == 2 {
                     for (index, name) in indices.iter() {
                         if *index == 0 {
@@ -4709,21 +4793,6 @@ impl<'a> CodeGenerator<'a> {
                         );
                     }
                 }
-
-                if complex_clause {
-                    let next_clause = arg_stack.pop().unwrap();
-
-                    term = apply_wrap(
-                        Term::Lambda {
-                            parameter_name: Name {
-                                text: "__other_clauses_delayed".to_string(),
-                                unique: 0.into(),
-                            },
-                            body: term.into(),
-                        },
-                        Term::Delay(next_clause.into()),
-                    )
-                }
                 arg_stack.push(term);
             }
         }