Fix couple of bugs in handling returns in if blocks

sway-core/src/ir_generation/compile.rs

@@ -230,6 +230,8 @@ fn compile_fn_with_args(
         ret_val = Constant::get_unit(context, None);
     }
 
+    let already_returns = compiler.current_block.is_terminated_by_ret(context);
+
     // Another special case: if the last expression in a function is a return then we don't want to
     // add another implicit return instruction here, as `ret_val` will be unit regardless of the
     // function return type actually is.  This new RET will be going into an unreachable block
@@ -239,9 +241,10 @@ fn compile_fn_with_args(
     // To tell if this is the case we can check that the current block is empty and has no
     // predecessors (and isn't the entry block which has none by definition), implying the most
     // recent instruction was a RET.
-    if compiler.current_block.num_instructions(context) > 0
-        || compiler.current_block == compiler.function.get_entry_block(context)
-        || compiler.current_block.num_predecessors(context) > 0
+    if !already_returns
+        && (compiler.current_block.num_instructions(context) > 0
+            || compiler.current_block == compiler.function.get_entry_block(context)
+            || compiler.current_block.num_predecessors(context) > 0)
     {
         if ret_type.eq(context, &Type::Unit) {
             ret_val = Constant::get_unit(context, None);

sway-core/src/ir_generation/function.rs

@@ -369,9 +369,6 @@ impl FnCompiler {
                 self.current_block
                     .ins(context)
                     .ret(ret_value, ret_ty, span_md_idx);
-                // RET is a terminator so we must create a new block here.  If anything is added to
-                // it then it'll almost certainly be dead code.
-                self.current_block = self.function.create_block(context, None);
                 Ok(Constant::get_unit(context, span_md_idx))
             }
         }
@@ -724,14 +721,15 @@ impl FnCompiler {
         if codeblock.contents.is_empty() {
             Some(insert_type(TypeInfo::Tuple(Vec::new())))
         } else {
-            codeblock
-                .contents
-                .iter()
-                .find_map(|node| match &node.content {
-                    TypedAstNodeContent::ReturnStatement(trs) => Some(trs.expr.return_type),
-                    TypedAstNodeContent::ImplicitReturnExpression(te) => Some(te.return_type),
-                    _otherwise => None,
-                })
+            codeblock.contents.iter().find_map(|node| {
+                match node.gather_return_statements().first() {
+                    Some(TypedReturnStatement { expr }) => Some(expr.return_type),
+                    None => match &node.content {
+                        TypedAstNodeContent::ImplicitReturnExpression(te) => Some(te.return_type),
+                        _otherwise => None,
+                    },
+                }
+            })
         }
     }
 
@@ -765,6 +763,7 @@ impl FnCompiler {
         self.current_block = true_block_begin;
         let true_value = self.compile_expression(context, ast_then)?;
         let true_block_end = self.current_block;
+        let then_returns = true_block_end.is_terminated_by_ret(context);
 
         let false_block_begin = self.function.create_block(context, None);
         self.current_block = false_block_begin;
@@ -773,6 +772,7 @@ impl FnCompiler {
             Some(expr) => self.compile_expression(context, *expr)?,
         };
         let false_block_end = self.current_block;
+        let else_returns = false_block_end.is_terminated_by_ret(context);
 
         entry_block.ins(context).conditional_branch(
             cond_value,
@@ -782,16 +782,28 @@ impl FnCompiler {
             cond_span_md_idx,
         );
 
+        if then_returns && else_returns {
+            return Ok(Constant::get_unit(context, None));
+        }
+
         let merge_block = self.function.create_block(context, None);
-        true_block_end
-            .ins(context)
-            .branch(merge_block, Some(true_value), None);
-        false_block_end
-            .ins(context)
-            .branch(merge_block, Some(false_value), None);
+        if !then_returns {
+            true_block_end
+                .ins(context)
+                .branch(merge_block, Some(true_value), None);
+        }
+        if !else_returns {
+            false_block_end
+                .ins(context)
+                .branch(merge_block, Some(false_value), None);
+        }
 
         self.current_block = merge_block;
-        Ok(merge_block.get_phi(context))
+        if !then_returns || !else_returns {
+            Ok(merge_block.get_phi(context))
+        } else {
+            Ok(Constant::get_unit(context, None))
+        }
     }
 
     fn compile_unsafe_downcast(

sway-core/src/semantic_analysis/ast_node/code_block.rs

@@ -76,7 +76,7 @@ impl TypedCodeBlock {
                         ..
                     }),
                 ..
-            } => Some(*return_type),
+            } if !x.deterministically_aborts() => Some(*return_type),
             _ => None,
         });
 

sway-core/src/semantic_analysis/ast_node/expression/typed_expression/if_expression.rs

@@ -37,11 +37,16 @@ pub(crate) fn instantiate_if_expression(
     let mut else_deterministically_aborts = false;
     let r#else = r#else.map(|r#else| {
         else_deterministically_aborts = r#else.deterministically_aborts();
+        let ty_to_check = if then_deterministically_aborts {
+            type_annotation
+        } else {
+            then.return_type
+        };
         if !else_deterministically_aborts {
             // if this does not deterministically_abort, check the block return type
             let (mut new_warnings, new_errors) = unify_with_self(
                 r#else.return_type,
-                then.return_type,
+                ty_to_check,
                 self_type,
                 &r#else.span,
                 "`else` branch must return expected type.",
@@ -78,7 +83,11 @@ pub(crate) fn instantiate_if_expression(
         }
     }
 
-    let return_type = then.return_type;
+    let return_type = if !then_deterministically_aborts {
+        then.return_type
+    } else {
+        r#else_ret_ty
+    };
     let exp = TypedExpression {
         expression: TypedExpressionVariant::IfExp {
             condition: Box::new(condition),

sway-ir/src/block.rs

@@ -149,6 +149,11 @@ impl Block {
         })
     }
 
+    pub fn is_terminated_by_ret(&self, context: &Context) -> bool {
+        self.get_term_inst(context)
+            .map_or(false, |i| matches!(i, Instruction::Ret { .. }))
+    }
+
     /// Replace a value within this block.
     ///
     /// For every instruction within the block, any reference to `old_val` is replaced with

test/src/e2e_vm_tests/mod.rs

@@ -429,6 +429,7 @@ pub fn run(locked: bool, filter_regex: Option<regex::Regex>) {
             ProgramState::Revert(0), // false
         ),
         ("should_pass/language/const_inits", ProgramState::Return(1)),
+        ("should_pass/language/impure_ifs", ProgramState::Return(2)),
         (
             "should_pass/language/enum_padding",
             ProgramState::ReturnData(Bytes32::from([

test/src/e2e_vm_tests/test_programs/should_pass/language/impure_ifs/Forc.lock

@@ -0,0 +1,14 @@
+[[package]]
+name = 'const_inits'
+source = 'root'
+dependencies = ['std']
+
+[[package]]
+name = 'core'
+source = 'path+from-root-76FD265A93458D61'
+dependencies = []
+
+[[package]]
+name = 'std'
+source = 'path+from-root-76FD265A93458D61'
+dependencies = ['core']

test/src/e2e_vm_tests/test_programs/should_pass/language/impure_ifs/Forc.toml

@@ -0,0 +1,8 @@
+[project]
+authors = ["Fuel Labs <[email protected]>"]
+entry = "main.sw"
+license = "Apache-2.0"
+name = "const_inits"
+
+[dependencies]
+std = { path = "../../../../../../../sway-lib-std" }

test/src/e2e_vm_tests/test_programs/should_pass/language/impure_ifs/json_abi_oracle.json

@@ -0,0 +1,14 @@
+[
+  {
+    "inputs": [],
+    "name": "main",
+    "outputs": [
+      {
+        "components": null,
+        "name": "",
+        "type": "u64"
+      }
+    ],
+    "type": "function"
+  }
+]

test/src/e2e_vm_tests/test_programs/should_pass/language/impure_ifs/src/main.sw

@@ -0,0 +1,90 @@
+script;
+
+use std::{assert::assert, logging::log};
+
+enum Bool {
+    True: (),
+    False: (),
+}
+
+fn foo(b: bool) -> u64 {
+    if b {
+        101
+    } else {
+        102
+    }
+}
+
+fn bar(b: bool) -> u64 {
+    if b {
+        return 101;
+    } else {
+        return 102;
+    }
+}
+
+fn bell(b: bool) -> u64 {
+    if b {
+        return 101;
+    } else {
+        102
+    }
+}
+
+fn moo(b: bool) -> u64 {
+    if b {
+        101
+    } else {
+        return 102;
+    }
+}
+
+fn poo(b: Bool) -> u64 {
+    if let Bool::True = b {
+        101
+    } else {
+        return 102;
+    }
+}
+
+fn ran_out_of_names(b: Bool) -> u64 {
+    if let Bool::True = b {
+        return 101;
+    } else {
+        return 102;
+    }
+}
+
+fn another_fn(b: Bool) -> u64 {
+    if let Bool::True = b {
+        return 101;
+    } else {
+        102
+    }
+}
+
+fn thats_all(b: Bool) -> u64 {
+    if let Bool::True = b {
+        101
+    } else {
+        102
+    }
+}
+
+fn main() -> u64 {
+    assert(foo(true) == bar(true));
+    assert(foo(false) == bar(false));
+    assert(foo(true) == bell(true));
+    assert(foo(false) == bell(false));
+    assert(foo(true) == moo(true));
+    assert(foo(false) == moo(false));
+
+    assert(thats_all(Bool::True) == poo(Bool::True));
+    assert(thats_all(Bool::False) == poo(Bool::False));
+    assert(thats_all(Bool::True) == ran_out_of_names(Bool::True));
+    assert(thats_all(Bool::False) == ran_out_of_names(Bool::False));
+    assert(thats_all(Bool::True) == another_fn(Bool::True));
+    assert(thats_all(Bool::False) == another_fn(Bool::False));
+
+    2
+}