Merge pull request #109 from lambdaclass/fix-return-bug

Fix #107 Control Flow Bug

crates/concrete_driver/tests/examples.rs

@@ -16,6 +16,8 @@ mod common;
 #[test_case(include_str!("../../../examples/structs.con"), "structs", false, 8 ; "structs.con")]
 #[test_case(include_str!("../../../examples/casts.con"), "casts", false, 2 ; "casts.con")]
 #[test_case(include_str!("../../../examples/malloc.con"), "malloc", false, 5 ; "malloc.con")]
+#[test_case(include_str!("../../../examples/while_if_false.con"), "while_if_false", false, 7 ; "while_if_false.con")]
+#[test_case(include_str!("../../../examples/if_if_false.con"), "if_if_false", false, 7 ; "if_if_false.con")]
 fn example_tests(source: &str, name: &str, is_library: bool, status_code: i32) {
     assert_eq!(
         status_code,

crates/concrete_ir/src/lowering.rs

@@ -309,16 +309,14 @@ fn lower_func(
         lower_statement(&mut builder, stmt, ret_type.clone())?;
     }
 
-    if !builder.statements.is_empty() {
-        let statements = std::mem::take(&mut builder.statements);
-        builder.body.basic_blocks.push(BasicBlock {
-            statements,
-            terminator: Box::new(Terminator {
-                span: None,
-                kind: TerminatorKind::Return,
-            }),
-        });
-    }
+    let statements = std::mem::take(&mut builder.statements);
+    builder.body.basic_blocks.push(BasicBlock {
+        statements,
+        terminator: Box::new(Terminator {
+            span: None,
+            kind: TerminatorKind::Return,
+        }),
+    });
 
     let (mut ctx, body) = (builder.ctx, builder.body);
     ctx.unresolved_function_signatures.remove(&body.id);
@@ -366,12 +364,18 @@ fn lower_statement(
         statements::Statement::Assign(info) => lower_assign(builder, info)?,
         statements::Statement::Match(_) => todo!(),
         statements::Statement::For(_) => todo!(),
-        statements::Statement::If(info) => lower_if_statement(builder, info)?,
+        statements::Statement::If(info) => {
+            lower_if_statement(builder, info)?;
+            assert!(builder.statements.is_empty());
+        }
         statements::Statement::Let(info) => lower_let(builder, info)?,
         statements::Statement::Return(info) => {
             lower_return(builder, info, ret_type)?;
         }
-        statements::Statement::While(info) => lower_while(builder, info)?,
+        statements::Statement::While(info) => {
+            lower_while(builder, info)?;
+            assert!(builder.statements.is_empty());
+        }
         statements::Statement::FnCall(info) => {
             lower_fn_call(builder, info)?;
         }
@@ -428,25 +432,17 @@ fn lower_while(builder: &mut FnBodyBuilder, info: &WhileStmt) -> Result<(), Lowe
         )?;
     }
 
-    // keet idx to change terminator if there is no return
-    let last_then_block_idx = if !matches!(
-        builder.body.basic_blocks.last().unwrap().terminator.kind,
-        TerminatorKind::Return
-    ) {
-        builder.body.basic_blocks.len();
-        let statements = std::mem::take(&mut builder.statements);
-        let idx = builder.body.basic_blocks.len();
-        builder.body.basic_blocks.push(BasicBlock {
-            statements,
-            terminator: Box::new(Terminator {
-                span: None,
-                kind: TerminatorKind::Unreachable,
-            }),
-        });
-        Some(idx)
-    } else {
-        None
-    };
+    builder.body.basic_blocks.len();
+    let statements = std::mem::take(&mut builder.statements);
+    builder.body.basic_blocks.push(BasicBlock {
+        statements,
+        terminator: Box::new(Terminator {
+            span: None,
+            kind: TerminatorKind::Goto {
+                target: check_block_idx,
+            },
+        }),
+    });
 
     let otherwise_block_idx = builder.body.basic_blocks.len();
 
@@ -461,14 +457,6 @@ fn lower_while(builder: &mut FnBodyBuilder, info: &WhileStmt) -> Result<(), Lowe
     };
     builder.body.basic_blocks[check_block_idx].terminator.kind = kind;
 
-    if let Some(last_then_block_idx) = last_then_block_idx {
-        builder.body.basic_blocks[last_then_block_idx]
-            .terminator
-            .kind = TerminatorKind::Goto {
-            target: check_block_idx,
-        };
-    }
-
     Ok(())
 }
 
@@ -511,23 +499,17 @@ fn lower_if_statement(builder: &mut FnBodyBuilder, info: &IfExpr) -> Result<(),
     }
 
     // keet idx to change terminator
-    let last_then_block_idx = if !matches!(
-        builder.body.basic_blocks.last().unwrap().terminator.kind,
-        TerminatorKind::Return
-    ) {
+    let last_then_block_idx = {
         builder.body.basic_blocks.len();
         let statements = std::mem::take(&mut builder.statements);
-        let idx = builder.body.basic_blocks.len();
         builder.body.basic_blocks.push(BasicBlock {
             statements,
             terminator: Box::new(Terminator {
                 span: None,
                 kind: TerminatorKind::Unreachable,
             }),
         });
-        Some(idx)
-    } else {
-        None
+        builder.body.basic_blocks.len() - 1
     };
 
     let first_else_block_idx = builder.body.basic_blocks.len();
@@ -540,32 +522,23 @@ fn lower_if_statement(builder: &mut FnBodyBuilder, info: &IfExpr) -> Result<(),
                 builder.body.locals[builder.ret_local].ty.clone(),
             )?;
         }
-    }
 
-    let last_else_block_idx = if !matches!(
-        builder.body.basic_blocks.last().unwrap().terminator.kind,
-        TerminatorKind::Return
-    ) {
-        builder.body.basic_blocks.len();
         let statements = std::mem::take(&mut builder.statements);
-        let idx = builder.body.basic_blocks.len();
         builder.body.basic_blocks.push(BasicBlock {
             statements,
             terminator: Box::new(Terminator {
                 span: None,
-                kind: TerminatorKind::Unreachable,
+                kind: TerminatorKind::Goto {
+                    target: builder.body.basic_blocks.len() + 1,
+                },
             }),
         });
-        Some(idx)
-    } else {
-        None
-    };
+    }
 
     let targets = SwitchTargets {
         values: vec![discriminator_type.kind.get_falsy_value()],
         targets: vec![first_else_block_idx, first_then_block_idx],
     };
-
     let kind = TerminatorKind::SwitchInt {
         discriminator: Operand::Place(place),
         targets,
@@ -574,22 +547,11 @@ fn lower_if_statement(builder: &mut FnBodyBuilder, info: &IfExpr) -> Result<(),
 
     let next_block_idx = builder.body.basic_blocks.len();
 
-    // check if the
-    if let Some(last_then_block_idx) = last_then_block_idx {
-        builder.body.basic_blocks[last_then_block_idx]
-            .terminator
-            .kind = TerminatorKind::Goto {
-            target: next_block_idx,
-        };
-    }
-
-    if let Some(last_else_block_idx) = last_else_block_idx {
-        builder.body.basic_blocks[last_else_block_idx]
-            .terminator
-            .kind = TerminatorKind::Goto {
-            target: next_block_idx,
-        };
-    }
+    builder.body.basic_blocks[last_then_block_idx]
+        .terminator
+        .kind = TerminatorKind::Goto {
+        target: next_block_idx,
+    };
 
     Ok(())
 }

examples/if_if_false.con

@@ -0,0 +1,15 @@
+mod Example {
+    fn main() -> i32 {
+        let foo: i32 = 7;
+
+        if true {
+            if false {
+                return 1;
+            }
+        } else {
+            foo = 14;
+        }
+
+        return foo;
+    }
+}

examples/while_if_false.con

@@ -0,0 +1,14 @@
+mod Example {
+    fn main() -> i32 {
+        let mut result: i32 = 7;
+
+        while false {
+            if false {
+                return 0;
+            }
+            result = 14;
+        }
+
+        return result;
+    }
+}