[BOLT] Fix X86MCPlusBuilder::replaceRegWithImm

Reassigning the operand didn't update the operand type which resulted in an
assertion (`Assertion `isReg() && "This is not a register operand!"' failed.`)
Reset the instruction instead.

Test Plan:
```
ninja check-bolt
...
PASS: BOLT-Unit :: Core/./CoreTests/X86/MCPlusBuilderTester.ReplaceRegWithImm/0 (90 of 136)
```

Reviewed By: rafauler

Differential Revision: https://reviews.llvm.org/D120263

3 files changed, 31 insertions, 16 deletions

diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index ac4416da7cc4..3a1492047cb2 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -3030,21 +3030,13 @@ public:
     if (NumFound != 1)
       return false;
 
-    // Iterate backwards to replace the src register before the src/dest
-    // register as in AND, ADD, and SUB Only iterate through src operands that
-    // arent also dest operands
-    for (unsigned Index = InstDesc.getNumOperands() - 1,
-                  E = InstDesc.getNumDefs() + (I.HasLHS ? 0 : -1);
-         Index != E; --Index) {
-      if (!Inst.getOperand(Index).isReg() ||
-          Inst.getOperand(Index).getReg() != Register)
-        continue;
-      MCOperand NewOperand = MCOperand::createImm(Imm);
-      Inst.getOperand(Index) = NewOperand;
-      break;
-    }
-
+    MCOperand TargetOp = Inst.getOperand(0);
+    Inst.clear();
     Inst.setOpcode(NewOpcode);
+    Inst.addOperand(TargetOp);
+    if (I.HasLHS)
+      Inst.addOperand(TargetOp);
+    Inst.addOperand(MCOperand::createImm(Imm));
 
     return true;
   }
diff --git a/bolt/test/X86/tail-duplication-prop-bug.s b/bolt/test/X86/tail-duplication-prop-bug.s
index b6dcbe70a35b..b91e85ddd27a 100644
--- a/bolt/test/X86/tail-duplication-prop-bug.s
+++ b/bolt/test/X86/tail-duplication-prop-bug.s
@@ -1,15 +1,16 @@
 # This reproduces a bug in aggressive tail duplication/copy propagation.
-# XFAIL: *
 
 # REQUIRES: system-linux
 # RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
 # RUN: link_fdata %s %t.o %t.fdata
 # RUN: llvm-strip --strip-unneeded %t.o
-# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -nostdlib
+# RUN: ld.lld %t.o -o %t.exe -q -nostdlib
 # RUN: llvm-bolt %t.exe -o %t.out -data %t.fdata -relocs \
 # RUN:   -tail-duplication=1 -tail-duplication-aggressive=1 \
 # RUN:   -tail-duplication-const-copy-propagation=1
 
+  .text
+  .type a, %function
   .globl a
 a:
 	.cfi_startproc
@@ -33,7 +34,10 @@ i:
 g:
 	jmp	g
 j:
+    ud2
 	.cfi_endproc
+.size a, .-a
+
 .rodata
 JT:
 	.quad	b
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index a1ebde8c1123..57326125e351 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -6,6 +6,8 @@
 #include "X86Subtarget.h"
 #endif // X86_AVAILABLE
 
+#include "bolt/Core/BinaryBasicBlock.h"
+#include "bolt/Core/BinaryFunction.h"
 #include "bolt/Rewrite/RewriteInstance.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
@@ -109,6 +111,23 @@ TEST_P(MCPlusBuilderTester, AliasSmallerAX) {
   testRegAliases(Triple::x86_64, X86::AX, AliasesAX, AliasesAXCount, true);
 }
 
+TEST_P(MCPlusBuilderTester, ReplaceRegWithImm) {
+  if (GetParam() != Triple::x86_64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock(0);
+  MCInst Inst; // cmpl    %eax, %ebx
+  Inst.setOpcode(X86::CMP32rr);
+  Inst.addOperand(MCOperand::createReg(X86::EAX));
+  Inst.addOperand(MCOperand::createReg(X86::EBX));
+  auto II = BB->addInstruction(Inst);
+  bool Replaced = BC->MIB->replaceRegWithImm(*II, X86::EBX, 1);
+  ASSERT_TRUE(Replaced);
+  ASSERT_EQ(II->getOpcode(), X86::CMP32ri8);
+  ASSERT_EQ(II->getOperand(0).getReg(), X86::EAX);
+  ASSERT_EQ(II->getOperand(1).getImm(), 1);
+}
+
 #endif // X86_AVAILABLE
 
 TEST_P(MCPlusBuilderTester, Annotation) {