Lower BoolOperation expressions

Fixes #488

crates/lowering/src/ast_utils.rs

@@ -1,5 +1,5 @@
 use fe_analyzer::namespace::types::FixedSize;
-use fe_parser::ast::{CallArg, Expr, FuncStmt, VarDeclTarget};
+use fe_parser::ast::{BoolOperator, CallArg, Expr, FuncStmt, UnaryOperator, VarDeclTarget};
 use fe_parser::node::{Node, NodeId};
 
 use crate::names;
@@ -494,6 +494,85 @@ pub fn ternary_to_if(
     unreachable!()
 }
 
+/// Turns a boolean expression into a set of statements resembling an if/else block with equal
+/// functionality. Expects the type and variable result name to be provided as parameters.
+pub fn boolean_expr_to_if(
+    expr_type: FixedSize,
+    expr: &Node<Expr>,
+    result_name: &str,
+) -> Vec<Node<FuncStmt>> {
+    if let Expr::BoolOperation { left, op, right } = &expr.kind {
+        return match op.kind {
+            BoolOperator::And => {
+                // from: left && right
+
+                // into:
+                // res: bool = false
+                // if left:
+                //     res = right
+
+                let mut stmts = vec![FuncStmt::VarDecl {
+                    target: VarDeclTarget::Name(result_name.to_string()).into_node(),
+                    typ: names::fixed_size_type_desc(&expr_type).into_node(),
+                    value: Some(Expr::Bool(false).into_node()),
+                }
+                .into_node()];
+                let if_branch = FuncStmt::Assign {
+                    target: Expr::Name(result_name.to_string()).into_node(),
+                    value: right.kind.clone().into_traceable_node(right.original_id),
+                }
+                .into_node();
+
+                stmts.push(
+                    FuncStmt::If {
+                        test: left.kind.clone().into_traceable_node(left.original_id),
+                        body: vec![if_branch],
+                        or_else: vec![],
+                    }
+                    .into_node(),
+                );
+                stmts
+            }
+            BoolOperator::Or => {
+                // from: left || right
+                // into:
+                // res: bool = true
+                // if not left:
+                //     res = right
+                let mut stmts = vec![FuncStmt::VarDecl {
+                    target: VarDeclTarget::Name(result_name.to_string()).into_node(),
+                    typ: names::fixed_size_type_desc(&expr_type).into_node(),
+                    value: Some(Expr::Bool(true).into_node()),
+                }
+                .into_node()];
+                let if_branch = FuncStmt::Assign {
+                    target: Expr::Name(result_name.to_string()).into_node(),
+                    value: right.kind.clone().into_traceable_node(right.original_id),
+                }
+                .into_node();
+
+                stmts.push(
+                    FuncStmt::If {
+                        test: Expr::UnaryOperation {
+                            op: UnaryOperator::Not.into_node(),
+                            operand: Box::new(
+                                left.kind.clone().into_traceable_node(left.original_id),
+                            ),
+                        }
+                        .into_node(),
+                        body: vec![if_branch],
+                        or_else: vec![],
+                    }
+                    .into_node(),
+                );
+                stmts
+            }
+        };
+    }
+
+    unreachable!()
+}
+
 /// Returns a vector of expressions with all ternary expressions that are
 /// contained within the given function statement. The last expression
 /// in the list is the outermost ternary expression found in the statement.
@@ -524,6 +603,36 @@ pub fn get_first_ternary_expressions(nodes: &[Node<FuncStmt>]) -> Vec<Node<Expr>
     vec![]
 }
 
+/// Returns a vector of expressions with all boolean expressions that are
+/// contained within the given function statement. The last expression
+/// in the list is the outermost boolean expression found in the statement.
+pub fn get_all_boolean_expressions(node: &Node<FuncStmt>) -> Vec<Node<Expr>> {
+    let mut expressions = vec![];
+    map_ast_node(node.clone().into(), &mut |exp| {
+        if let StmtOrExpr::Expr(expr) = &exp {
+            if let Expr::BoolOperation { .. } = expr.kind {
+                expressions.push(expr.clone())
+            }
+        }
+
+        exp
+    });
+
+    expressions
+}
+
+/// For a given set of nodes returns the first set of boolean expressions that can be found.
+/// The last expression in the list is the outermost boolean expression found in the statement.
+pub fn get_first_boolean_expressions(nodes: &[Node<FuncStmt>]) -> Vec<Node<Expr>> {
+    for node in nodes {
+        let result = get_all_boolean_expressions(node);
+        if !result.is_empty() {
+            return result;
+        }
+    }
+    vec![]
+}
+
 /// In a given set of `nodes replaces a node that matches the `node_id` with a name expression node.
 pub fn replace_node_with_name_expression(
     nodes: &[Node<FuncStmt>],

crates/lowering/src/mappers/functions.rs

@@ -1,4 +1,6 @@
-use crate::ast_utils::get_first_ternary_expressions;
+use crate::ast_utils::{
+    boolean_expr_to_if, get_first_boolean_expressions, get_first_ternary_expressions,
+};
 use crate::ast_utils::{
     inject_before_expression, replace_node_with_name_expression, ternary_to_if,
 };
@@ -11,7 +13,7 @@ use crate::utils::ZeroSpanNode;
 use fe_analyzer::namespace::items::FunctionId;
 use fe_analyzer::namespace::types::{Base, Type};
 use fe_analyzer::namespace::types::{FixedSize, TypeDowncast};
-use fe_parser::ast::{self as fe, FuncStmt, RegularFunctionArg};
+use fe_parser::ast::{self as fe, Expr, FuncStmt, RegularFunctionArg};
 use fe_parser::node::Node;
 
 /// Lowers a function definition.
@@ -49,7 +51,20 @@ pub fn func_def(context: &mut ModuleContext, function: FunctionId) -> Node<fe::F
         lowered_body
     };
 
-    let lowered_body = lower_ternary_expressions(&mut fn_ctx, lowered_body);
+    let lowered_body = lower_iteratively(
+        &mut fn_ctx,
+        lowered_body,
+        "ternary_result",
+        &get_first_ternary_expressions,
+        &ternary_to_if,
+    );
+    let lowered_body = lower_iteratively(
+        &mut fn_ctx,
+        lowered_body,
+        "boolean_expr_result",
+        &get_first_boolean_expressions,
+        &boolean_expr_to_if,
+    );
 
     let param_types = {
         let params = &signature.params;
@@ -106,33 +121,37 @@ pub fn func_def(context: &mut ModuleContext, function: FunctionId) -> Node<fe::F
     Node::new(lowered_function, node.span)
 }
 
-fn lower_ternary_expressions(
+fn lower_iteratively(
     context: &mut FnContext,
     statements: Vec<Node<FuncStmt>>,
+    result_name: &str,
+    getter_fn: &dyn Fn(&[Node<FuncStmt>]) -> Vec<Node<Expr>>,
+    mapper_fn: &dyn Fn(FixedSize, &Node<Expr>, &str) -> Vec<Node<FuncStmt>>,
 ) -> Vec<Node<FuncStmt>> {
     let mut current_statements = statements;
 
     loop {
-        let terns = get_first_ternary_expressions(&current_statements);
+        let expressions = getter_fn(&current_statements);
 
-        if let Some(ternary) = terns.last() {
+        if let Some(current_expression) = expressions.last() {
             let expr_attr = context
-                .expression_attributes(ternary.original_id)
+                .expression_attributes(current_expression.original_id)
                 .expect("missing attributes");
 
-            let ternary_type =
+            let expression_type =
                 FixedSize::try_from(expr_attr.typ.clone()).expect("Not a fixed size");
 
-            let unique_name = context.make_unique_name("ternary_result");
-            let generated_if_else = ternary_to_if(ternary_type.clone(), ternary, &unique_name);
+            let unique_name = context.make_unique_name(result_name);
+            let generated_statements =
+                mapper_fn(expression_type.clone(), current_expression, &unique_name);
             current_statements = inject_before_expression(
                 &current_statements,
-                ternary.original_id,
-                &generated_if_else,
+                current_expression.original_id,
+                &generated_statements,
             );
             current_statements = replace_node_with_name_expression(
                 &current_statements,
-                ternary.original_id,
+                current_expression.original_id,
                 &unique_name,
             );
         } else {

crates/lowering/tests/lowering.rs

@@ -75,5 +75,6 @@ test_file! { module_const, "lowering/module_const.fe" }
 test_file! { module_fn, "lowering/module_fn.fe" }
 test_file! { struct_fn, "lowering/struct_fn.fe" }
 test_file! { ternary, "lowering/ternary.fe" }
+test_file! { and_or, "lowering/and_or.fe" }
 // TODO: the analyzer rejects lowered nested tuples.
 // test_file!(array_tuple, "lowering/array_tuple.fe");