From 6f0baea59416ad3172a91b3aab1d4ddb55e21969 Mon Sep 17 00:00:00 2001
From: Andreas Kling <kling@serenityos.org>
Date: Thu, 26 Oct 2023 15:35:25 +0200
Subject: [PATCH] LibJIT: Add jump_if_zero() and jump_if_not_zero() to
 Assembler

These can use test reg,reg on x86 which gives us a shorter encoding.
---
 Userland/Libraries/LibJIT/Assembler.h | 50 +++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/Userland/Libraries/LibJIT/Assembler.h b/Userland/Libraries/LibJIT/Assembler.h
index df8fe3c386..25f2565a11 100644
--- a/Userland/Libraries/LibJIT/Assembler.h
+++ b/Userland/Libraries/LibJIT/Assembler.h
@@ -281,8 +281,53 @@ struct Assembler {
         }
     }
 
+    void test(Operand lhs, Operand rhs)
+    {
+        if (lhs.type == Operand::Type::Reg && rhs.type == Operand::Type::Reg) {
+            emit8(0x48
+                | ((to_underlying(rhs.reg) >= 8) ? 1 << 2 : 0)
+                | ((to_underlying(lhs.reg) >= 8) ? 1 << 0 : 0));
+            emit8(0x85);
+            emit8(0xc0 | (encode_reg(rhs.reg) << 3) | encode_reg(lhs.reg));
+        } else if (lhs.type == Operand::Type::Reg && rhs.type == Operand::Type::Imm32) {
+            emit8(0x48 | ((to_underlying(lhs.reg) >= 8) ? 1 << 0 : 0));
+            emit8(0xf7);
+            emit8(0xc0 | encode_reg(lhs.reg));
+            emit32(rhs.offset_or_immediate);
+        } else {
+            VERIFY_NOT_REACHED();
+        }
+    }
+
+    void jump_if_zero(Operand reg, Label& label)
+    {
+        test(reg, reg);
+
+        // jz label (RIP-relative 32-bit offset)
+        emit8(0x0f);
+        emit8(0x84);
+        emit32(0xdeadbeef);
+        label.add_jump(m_output.size());
+    }
+
+    void jump_if_not_zero(Operand reg, Label& label)
+    {
+        test(reg, reg);
+
+        // jnz label (RIP-relative 32-bit offset)
+        emit8(0x0f);
+        emit8(0x85);
+        emit32(0xdeadbeef);
+        label.add_jump(m_output.size());
+    }
+
     void jump_if_equal(Operand lhs, Operand rhs, Label& label)
     {
+        if (rhs.type == Operand::Type::Imm32 && rhs.offset_or_immediate == 0) {
+            jump_if_zero(lhs, label);
+            return;
+        }
+
         cmp(lhs, rhs);
 
         // je label (RIP-relative 32-bit offset)
@@ -294,6 +339,11 @@ struct Assembler {
 
     void jump_if_not_equal(Operand lhs, Operand rhs, Label& label)
     {
+        if (rhs.type == Operand::Type::Imm32 && rhs.offset_or_immediate == 0) {
+            jump_if_not_zero(lhs, label);
+            return;
+        }
+
         cmp(lhs, rhs);
 
         // jne label (RIP-relative 32-bit offset)