YJIT: Branch directly when nil? is known from types

2 files changed, 34 insertions, 18 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 0551491bbb..1c48a1b040 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -3340,11 +3340,6 @@ fn gen_branchnil(
         gen_check_ints(asm, side_exit);
     }
 
-    // Test if the value is Qnil
-    // RUBY_Qnil    /* ...0000 1000 */
-    let val_opnd = ctx.stack_pop(1);
-    asm.cmp(val_opnd, Opnd::UImm(Qnil.into()));
-
     // Get the branch target instruction offsets
     let next_idx = jit_next_insn_idx(jit) as i32;
     let jump_idx = next_idx + jump_offset;
@@ -3357,18 +3352,29 @@ fn gen_branchnil(
         idx: jump_idx.try_into().unwrap(),
     };
 
-    // Generate the branch instructions
-    gen_branch(
-        jit,
-        ctx,
-        asm,
-        ocb,
-        jump_block,
-        ctx,
-        Some(next_block),
-        Some(ctx),
-        gen_branchnil_branch,
-    );
+    let val_type = ctx.get_opnd_type(StackOpnd(0));
+    let val_opnd = ctx.stack_pop(1);
+
+    if let Some(result) = val_type.known_nil() {
+        let target = if result { jump_block } else { next_block };
+        gen_direct_jump(jit, ctx, target, asm);
+    } else {
+        // Test if the value is Qnil
+        // RUBY_Qnil    /* ...0000 1000 */
+        asm.cmp(val_opnd, Opnd::UImm(Qnil.into()));
+        // Generate the branch instructions
+        gen_branch(
+            jit,
+            ctx,
+            asm,
+            ocb,
+            jump_block,
+            ctx,
+            Some(next_block),
+            Some(ctx),
+            gen_branchnil_branch,
+        );
+    }
 
     EndBlock
 }
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 687dc21013..102988db11 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -172,7 +172,7 @@ impl Type {
         }
     }
 
-    /// Returns an Option with the exact value if it is known, otherwise None
+    /// Returns an Option boolean representing whether the value is truthy if known, otherwise None
     pub fn known_truthy(&self) -> Option<bool> {
         match self {
             Type::Nil => Some(false),
@@ -183,6 +183,16 @@ impl Type {
         }
     }
 
+    /// Returns an Option boolean representing whether the value is equal to nil if known, otherwise None
+    pub fn known_nil(&self) -> Option<bool> {
+        match (self, self.known_truthy()) {
+            (Type::Nil, _) => Some(true),
+            (Type::False, _) => Some(false), // Qfalse is not nil
+            (_, Some(true))  => Some(false), // if truthy, can't be nil
+            (_, _) => None // otherwise unknown
+        }
+    }
+
     /// Compute a difference between two value types
     /// Returns 0 if the two are the same
     /// Returns > 0 if different but compatible