Fix couple of bugs in handling returns in if blocks

sway-core/src/ir_generation/function.rs

@@ -5,6 +5,7 @@ use crate::{
     parse_tree::{AsmOp, AsmRegister, LazyOp, Literal, Purity, Visibility},
     semantic_analysis::*,
     type_engine::{insert_type, resolve_type, TypeId, TypeInfo},
+    types::deterministically_aborts::DeterministicallyAborts,
 };
 
 use super::{compile::compile_function, convert::*, lexical_map::LexicalMap, types::*};
@@ -724,14 +725,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,
+                    },
+                }
+            })
         }
     }
 
@@ -748,6 +750,12 @@ impl FnCompiler {
         let cond_value = self.compile_expression(context, ast_condition)?;
         let entry_block = self.current_block;
 
+        let then_deterministically_aborts = ast_then.deterministically_aborts();
+        let else_deterministically_aborts = ast_else
+            .as_ref()
+            .map(|x| x.deterministically_aborts())
+            .unwrap_or(false);
+
         // To keep the blocks in a nice order we create them only as we populate them.  It's
         // possible when compiling other expressions for the 'current' block to change, and it
         // should always be the block to which instructions are added.  So for the true and false
@@ -782,16 +790,37 @@ impl FnCompiler {
             cond_span_md_idx,
         );
 
+        if then_deterministically_aborts && else_deterministically_aborts {
+            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);
+        true_block_end.ins(context).branch(
+            merge_block,
+            if !then_deterministically_aborts {
+                Some(true_value)
+            } else {
+                None
+            },
+            None,
+        );
+        false_block_end.ins(context).branch(
+            merge_block,
+            if !else_deterministically_aborts {
+                Some(false_value)
+            } else {
+                None
+            },
+            None,
+        );
 
         self.current_block = merge_block;
-        Ok(merge_block.get_phi(context))
+
+        if !then_deterministically_aborts || !else_deterministically_aborts {
+            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-core/src/types/deterministically_aborts.rs

@@ -1,5 +1,5 @@
 /// If this expression deterministically_aborts 100% of the time, this function returns
 /// `true`. Used in dead-code and control-flow analysis.
-pub(crate) trait DeterministicallyAborts {
+pub trait DeterministicallyAborts {
     fn deterministically_aborts(&self) -> bool;
 }

sway-core/src/types/mod.rs

@@ -1,7 +1,7 @@
-mod deterministically_aborts;
+pub mod deterministically_aborts;
 mod json_abi_string;
 mod to_json_abi;
 
-pub(crate) use deterministically_aborts::*;
+pub use deterministically_aborts::*;
 pub(crate) use json_abi_string::*;
 pub use to_json_abi::*;

sway-ir/src/constant.rs

@@ -171,4 +171,21 @@ impl Constant {
         ));
         Value::new_constant(context, value, span_md_idx)
     }
+
+    pub fn eq(&self, context: &Context, other: &Constant) -> bool {
+        self.ty.eq(context, &other.ty)
+            && match (&self.value, &other.value) {
+                (ConstantValue::Undef, ConstantValue::Undef)
+                | (ConstantValue::Unit, ConstantValue::Unit) => true,
+                (ConstantValue::Bool(v1), ConstantValue::Bool(v2)) => v1 == v2,
+                (ConstantValue::Uint(v1), ConstantValue::Uint(v2)) => v1 == v2,
+                (ConstantValue::B256(a1), ConstantValue::B256(a2)) => a1 == a2,
+                (ConstantValue::String(s1), ConstantValue::String(s2)) => s1 == s2,
+                (ConstantValue::Array(a1), ConstantValue::Array(a2))
+                | (ConstantValue::Struct(a1), ConstantValue::Struct(a2)) => {
+                    a1.iter().zip(a2.iter()).all(|(c1, c2)| c1.eq(context, c2))
+                }
+                _ => false,
+            }
+    }
 }

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,46 @@
+script;
+
+use std::{assert::assert, logging::log};
+
+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 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));
+
+    2
+}