[BOLT] Handle generation of compare and jump sequences

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -1714,6 +1714,15 @@ class MCPlusBuilder {
     return {};
   }
 
+  /// Create a sequence of instructions to compare contents of a register
+  /// \p RegNo to immediate \Imm and jump to \p Target if they are different.
+  virtual InstructionListType createCmpJNE(MCPhysReg RegNo, int64_t Imm,
+                                           const MCSymbol *Target,
+                                           MCContext *Ctx) const {
+    llvm_unreachable("not implemented");
+    return {};
+  }
+
   /// Creates inline memcpy instruction. If \p ReturnEnd is true, then return
   /// (dest + n) instead of dest.
   virtual InstructionListType createInlineMemcpy(bool ReturnEnd) const {

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -1321,17 +1321,47 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
 
   int getUncondBranchEncodingSize() const override { return 28; }
 
+  // This helper function creates the snippet of code that compares a register
+  // RegNo with an immedaite Imm, and jumps to Target if they are equal.
+  // cmp RegNo, #Imm
+  // b.eq Target
+  // where cmp is an alias for subs, which results in the code below:
+  // subs xzr, RegNo, #Imm
+  // b.eq Target.
   InstructionListType createCmpJE(MCPhysReg RegNo, int64_t Imm,
                                   const MCSymbol *Target,
                                   MCContext *Ctx) const override {
     InstructionListType Code;
     Code.emplace_back(MCInstBuilder(AArch64::SUBSXri)
-                          .addReg(RegNo)
+                          .addReg(AArch64::XZR)
                           .addReg(RegNo)
                           .addImm(Imm)
                           .addImm(0));
     Code.emplace_back(MCInstBuilder(AArch64::Bcc)
+                          .addImm(AArch64CC::EQ)
+                          .addExpr(MCSymbolRefExpr::create(
+                              Target, MCSymbolRefExpr::VK_None, *Ctx)));
+    return Code;
+  }
+
+  // This helper function creates the snippet of code that compares a register
+  // RegNo with an immedaite Imm, and jumps to Target if they are not equal.
+  // cmp RegNo, #Imm
+  // b.ne Target
+  // where cmp is an alias for subs, which results in the code below:
+  // subs xzr, RegNo, #Imm
+  // b.ne Target.
+  InstructionListType createCmpJNE(MCPhysReg RegNo, int64_t Imm,
+                                   const MCSymbol *Target,
+                                   MCContext *Ctx) const override {
+    InstructionListType Code;
+    Code.emplace_back(MCInstBuilder(AArch64::SUBSXri)
+                          .addReg(AArch64::XZR)
+                          .addReg(RegNo)
                           .addImm(Imm)
+                          .addImm(0));
+    Code.emplace_back(MCInstBuilder(AArch64::Bcc)
+                          .addImm(AArch64CC::NE)
                           .addExpr(MCSymbolRefExpr::create(
                               Target, MCSymbolRefExpr::VK_None, *Ctx)));
     return Code;

bolt/lib/Target/X86/X86MCPlusBuilder.cpp

@@ -2436,6 +2436,18 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     return Code;
   }
 
+  InstructionListType createCmpJNE(MCPhysReg RegNo, int64_t Imm,
+                                   const MCSymbol *Target,
+                                   MCContext *Ctx) const override {
+    InstructionListType Code;
+    Code.emplace_back(MCInstBuilder(X86::CMP64ri8).addReg(RegNo).addImm(Imm));
+    Code.emplace_back(MCInstBuilder(X86::JCC_1)
+                          .addExpr(MCSymbolRefExpr::create(
+                              Target, MCSymbolRefExpr::VK_None, *Ctx))
+                          .addImm(X86::COND_NE));
+    return Code;
+  }
+
   std::optional<Relocation>
   createRelocation(const MCFixup &Fixup,
                    const MCAsmBackend &MAB) const override {

bolt/unittests/Core/MCPlusBuilder.cpp

@@ -107,6 +107,54 @@ TEST_P(MCPlusBuilderTester, AliasSmallerX0) {
   testRegAliases(Triple::aarch64, AArch64::X0, AliasesX0, AliasesX0Count, true);
 }
 
+TEST_P(MCPlusBuilderTester, AArch64_CmpJE) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+
+  InstructionListType Instrs =
+      BC->MIB->createCmpJE(AArch64::X0, 2, BB->getLabel(), BC->Ctx.get());
+  BB->addInstructions(Instrs.begin(), Instrs.end());
+  BB->addSuccessor(BB.get());
+
+  auto II = BB->begin();
+  ASSERT_EQ(II->getOpcode(), AArch64::SUBSXri);
+  ASSERT_EQ(II->getOperand(0).getReg(), AArch64::XZR);
+  ASSERT_EQ(II->getOperand(1).getReg(), AArch64::X0);
+  ASSERT_EQ(II->getOperand(2).getImm(), 2);
+  ASSERT_EQ(II->getOperand(3).getImm(), 0);
+  II++;
+  ASSERT_EQ(II->getOpcode(), AArch64::Bcc);
+  ASSERT_EQ(II->getOperand(0).getImm(), AArch64CC::EQ);
+  const MCSymbol *Label = BC->MIB->getTargetSymbol(*II, 1);
+  ASSERT_EQ(Label, BB->getLabel());
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_CmpJNE) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+
+  InstructionListType Instrs =
+      BC->MIB->createCmpJNE(AArch64::X0, 2, BB->getLabel(), BC->Ctx.get());
+  BB->addInstructions(Instrs.begin(), Instrs.end());
+  BB->addSuccessor(BB.get());
+
+  auto II = BB->begin();
+  ASSERT_EQ(II->getOpcode(), AArch64::SUBSXri);
+  ASSERT_EQ(II->getOperand(0).getReg(), AArch64::XZR);
+  ASSERT_EQ(II->getOperand(1).getReg(), AArch64::X0);
+  ASSERT_EQ(II->getOperand(2).getImm(), 2);
+  ASSERT_EQ(II->getOperand(3).getImm(), 0);
+  II++;
+  ASSERT_EQ(II->getOpcode(), AArch64::Bcc);
+  ASSERT_EQ(II->getOperand(0).getImm(), AArch64CC::NE);
+  const MCSymbol *Label = BC->MIB->getTargetSymbol(*II, 1);
+  ASSERT_EQ(Label, BB->getLabel());
+}
+
 #endif // AARCH64_AVAILABLE
 
 #ifdef X86_AVAILABLE
@@ -143,6 +191,50 @@ TEST_P(MCPlusBuilderTester, ReplaceRegWithImm) {
   ASSERT_EQ(II->getOperand(1).getImm(), 1);
 }
 
+TEST_P(MCPlusBuilderTester, X86_CmpJE) {
+  if (GetParam() != Triple::x86_64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+
+  InstructionListType Instrs =
+      BC->MIB->createCmpJE(X86::EAX, 2, BB->getLabel(), BC->Ctx.get());
+  BB->addInstructions(Instrs.begin(), Instrs.end());
+  BB->addSuccessor(BB.get());
+
+  auto II = BB->begin();
+  ASSERT_EQ(II->getOpcode(), X86::CMP64ri8);
+  ASSERT_EQ(II->getOperand(0).getReg(), X86::EAX);
+  ASSERT_EQ(II->getOperand(1).getImm(), 2);
+  II++;
+  ASSERT_EQ(II->getOpcode(), X86::JCC_1);
+  const MCSymbol *Label = BC->MIB->getTargetSymbol(*II, 0);
+  ASSERT_EQ(Label, BB->getLabel());
+  ASSERT_EQ(II->getOperand(1).getImm(), X86::COND_E);
+}
+
+TEST_P(MCPlusBuilderTester, X86_CmpJNE) {
+  if (GetParam() != Triple::x86_64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+
+  InstructionListType Instrs =
+      BC->MIB->createCmpJNE(X86::EAX, 2, BB->getLabel(), BC->Ctx.get());
+  BB->addInstructions(Instrs.begin(), Instrs.end());
+  BB->addSuccessor(BB.get());
+
+  auto II = BB->begin();
+  ASSERT_EQ(II->getOpcode(), X86::CMP64ri8);
+  ASSERT_EQ(II->getOperand(0).getReg(), X86::EAX);
+  ASSERT_EQ(II->getOperand(1).getImm(), 2);
+  II++;
+  ASSERT_EQ(II->getOpcode(), X86::JCC_1);
+  const MCSymbol *Label = BC->MIB->getTargetSymbol(*II, 0);
+  ASSERT_EQ(Label, BB->getLabel());
+  ASSERT_EQ(II->getOperand(1).getImm(), X86::COND_NE);
+}
+
 #endif // X86_AVAILABLE
 
 TEST_P(MCPlusBuilderTester, Annotation) {