Merge pull request #1350 from rtfeldman/constrain_if

Constrain Expr2::If

editor/src/lang/ast.rs

@@ -119,10 +119,10 @@ pub enum Expr2 {
         elems: PoolVec<Expr2>, // 8B
     },
     If {
-        cond_var: Variable,                // 4B
-        expr_var: Variable,                // 4B
-        branches: PoolVec<(Expr2, Expr2)>, // 8B
-        final_else: NodeId<Expr2>,         // 4B
+        cond_var: Variable,                                // 4B
+        expr_var: Variable,                                // 4B
+        branches: PoolVec<(NodeId<Expr2>, NodeId<Expr2>)>, // 8B
+        final_else: NodeId<Expr2>,                         // 4B
     },
     When {
         cond_var: Variable,            // 4B

editor/src/lang/constrain.rs

@@ -13,7 +13,7 @@ use roc_module::{ident::TagName, symbol::Symbol};
 use roc_region::all::{Located, Region};
 use roc_types::{
     subs::Variable,
-    types,
+    types::{self, AnnotationSource},
     types::{Category, Reason},
 };
 
@@ -479,6 +479,157 @@ pub fn constrain_expr<'a>(
 
             exists(arena, flex_vars, And(and_constraints))
         }
+        Expr2::If {
+            cond_var,
+            expr_var,
+            branches,
+            final_else,
+        } => {
+            let expect_bool = |region| {
+                let bool_type = Type2::Variable(Variable::BOOL);
+                Expected::ForReason(Reason::IfCondition, bool_type, region)
+            };
+
+            let mut branch_cons = BumpVec::with_capacity_in(2 * branches.len() + 3, arena);
+
+            // TODO why does this cond var exist? is it for error messages?
+            // let first_cond_region = branches[0].0.region;
+            let cond_var_is_bool_con = Eq(
+                Type2::Variable(*cond_var),
+                expect_bool(region),
+                Category::If,
+                region,
+            );
+
+            branch_cons.push(cond_var_is_bool_con);
+
+            let final_else_expr = env.pool.get(*final_else);
+
+            let mut flex_vars = BumpVec::with_capacity_in(2, arena);
+
+            flex_vars.push(*cond_var);
+            flex_vars.push(*expr_var);
+
+            match expected {
+                Expected::FromAnnotation(name, arity, _, tipe) => {
+                    let num_branches = branches.len() + 1;
+
+                    for (index, branch_id) in branches.iter_node_ids().enumerate() {
+                        let (cond_id, body_id) = env.pool.get(branch_id);
+
+                        let cond = env.pool.get(*cond_id);
+                        let body = env.pool.get(*body_id);
+
+                        let cond_con =
+                            constrain_expr(arena, env, cond, expect_bool(region), region);
+
+                        let then_con = constrain_expr(
+                            arena,
+                            env,
+                            body,
+                            Expected::FromAnnotation(
+                                name.clone(),
+                                arity,
+                                AnnotationSource::TypedIfBranch {
+                                    index: Index::zero_based(index),
+                                    num_branches,
+                                },
+                                tipe.shallow_clone(),
+                            ),
+                            region,
+                        );
+
+                        branch_cons.push(cond_con);
+                        branch_cons.push(then_con);
+                    }
+
+                    let else_con = constrain_expr(
+                        arena,
+                        env,
+                        final_else_expr,
+                        Expected::FromAnnotation(
+                            name,
+                            arity,
+                            AnnotationSource::TypedIfBranch {
+                                index: Index::zero_based(branches.len()),
+                                num_branches,
+                            },
+                            tipe.shallow_clone(),
+                        ),
+                        region,
+                    );
+
+                    let ast_con = Eq(
+                        Type2::Variable(*expr_var),
+                        Expected::NoExpectation(tipe),
+                        Category::Storage(std::file!(), std::line!()),
+                        region,
+                    );
+
+                    branch_cons.push(ast_con);
+                    branch_cons.push(else_con);
+
+                    exists(arena, flex_vars, And(branch_cons))
+                }
+                _ => {
+                    for (index, branch_id) in branches.iter_node_ids().enumerate() {
+                        let (cond_id, body_id) = env.pool.get(branch_id);
+
+                        let cond = env.pool.get(*cond_id);
+                        let body = env.pool.get(*body_id);
+
+                        let cond_con =
+                            constrain_expr(arena, env, cond, expect_bool(region), region);
+
+                        let then_con = constrain_expr(
+                            arena,
+                            env,
+                            body,
+                            Expected::ForReason(
+                                Reason::IfBranch {
+                                    index: Index::zero_based(index),
+                                    total_branches: branches.len(),
+                                },
+                                Type2::Variable(*expr_var),
+                                // should be from body
+                                region,
+                            ),
+                            region,
+                        );
+
+                        branch_cons.push(cond_con);
+                        branch_cons.push(then_con);
+                    }
+
+                    let else_con = constrain_expr(
+                        arena,
+                        env,
+                        final_else_expr,
+                        Expected::ForReason(
+                            Reason::IfBranch {
+                                index: Index::zero_based(branches.len()),
+                                total_branches: branches.len() + 1,
+                            },
+                            Type2::Variable(*expr_var),
+                            // should come from final_else
+                            region,
+                        ),
+                        region,
+                    );
+
+                    branch_cons.push(Eq(
+                        Type2::Variable(*expr_var),
+                        expected,
+                        Category::Storage(std::file!(), std::line!()),
+                        region,
+                    ));
+
+                    branch_cons.push(else_con);
+
+                    exists(arena, flex_vars, And(branch_cons))
+                }
+            }
+        }
         _ => todo!("implement constaints for {:?}", expr),
     }
 }

editor/src/lang/expr.rs

@@ -516,7 +516,7 @@ pub fn to_expr2<'a>(
                 output.references.union_mut(cond_output.references);
                 output.references.union_mut(then_output.references);
 
-                new_branches.push((cond, then_expr));
+                new_branches.push((env.pool.add(cond), env.pool.add(then_expr)));
             }
 
             let (else_expr, else_output) =

editor/tests/solve_expr2.rs

@@ -274,3 +274,15 @@ fn constrain_access() {
         "Str",
     )
 }
+
+#[test]
+fn constrain_if() {
+    infer_eq(
+        indoc!(
+            r#"
+            if True then Green else Red
+            "#
+        ),
+        "[ Green, Red ]*",
+    )
+}