Optimize branches when the target is statically known to a jump

This can happen in generic code

Author: bjorn3

src/base.rs

@@ -307,7 +307,9 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, impl Module>) {
             } => {
                 let discr = codegen_operand(fx, discr).load_scalar(fx);
 
-                if switch_ty.kind() == fx.tcx.types.bool.kind() {
+                let use_bool_opt = switch_ty.kind() == fx.tcx.types.bool.kind()
+                    || (targets.iter().count() == 1 && targets.iter().next().unwrap().0 == 0);
+                if use_bool_opt {
                     assert_eq!(targets.iter().count(), 1);
                     let (then_value, then_block) = targets.iter().next().unwrap();
                     let then_block = fx.get_block(then_block);
@@ -325,12 +327,22 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, impl Module>) {
                     let discr = crate::optimize::peephole::maybe_unwrap_bint(&mut fx.bcx, discr);
                     let discr =
                         crate::optimize::peephole::make_branchable_value(&mut fx.bcx, discr);
-                    if test_zero {
-                        fx.bcx.ins().brz(discr, then_block, &[]);
-                        fx.bcx.ins().jump(else_block, &[]);
+                    if let Some(taken) = crate::optimize::peephole::maybe_known_branch_taken(
+                        &fx.bcx, discr, test_zero,
+                    ) {
+                        if taken {
+                            fx.bcx.ins().jump(then_block, &[]);
+                        } else {
+                            fx.bcx.ins().jump(else_block, &[]);
+                        }
                     } else {
-                        fx.bcx.ins().brnz(discr, then_block, &[]);
-                        fx.bcx.ins().jump(else_block, &[]);
+                        if test_zero {
+                            fx.bcx.ins().brz(discr, then_block, &[]);
+                            fx.bcx.ins().jump(else_block, &[]);
+                        } else {
+                            fx.bcx.ins().brnz(discr, then_block, &[]);
+                            fx.bcx.ins().jump(else_block, &[]);
+                        }
                     }
                 } else {
                     let mut switch = ::cranelift_frontend::Switch::new();

src/optimize/peephole.rs

@@ -73,11 +73,48 @@ pub(crate) fn make_branchable_value(bcx: &mut FunctionBuilder<'_>, arg: Value) -
     })()
     .unwrap_or_else(|| {
         match bcx.func.dfg.value_type(arg) {
-            types::I8 | types::I32 => {
+            types::I8 | types::I16 => {
                 // WORKAROUND for brz.i8 and brnz.i8 not yet being implemented
                 bcx.ins().uextend(types::I32, arg)
             }
             _ => arg,
         }
     })
 }
+
+/// Returns whether the branch is statically known to be taken or `None` if it isn't statically known.
+pub(crate) fn maybe_known_branch_taken(
+    bcx: &FunctionBuilder<'_>,
+    arg: Value,
+    test_zero: bool,
+) -> Option<bool> {
+    let arg_inst = if let ValueDef::Result(arg_inst, 0) = bcx.func.dfg.value_def(arg) {
+        arg_inst
+    } else {
+        return None;
+    };
+
+    match bcx.func.dfg[arg_inst] {
+        InstructionData::UnaryBool {
+            opcode: Opcode::Bconst,
+            imm,
+        } => {
+            if test_zero {
+                Some(!imm)
+            } else {
+                Some(imm)
+            }
+        }
+        InstructionData::UnaryImm {
+            opcode: Opcode::Iconst,
+            imm,
+        } => {
+            if test_zero {
+                Some(imm.bits() == 0)
+            } else {
+                Some(imm.bits() != 0)
+            }
+        }
+        _ => None,
+    }
+}