[BOLT][RISCV] Carry-over annotations when fixing calls (#66763)

`FixRISCVCallsPass` changes all different forms of calls to `PseudoCALL`
instructions. However, the original call's annotations were lost in the
process.

This patch fixes this by moving all annotations from the old to the new
call. `MCPlusBuilder::moveAnnotations` had to be made public for this.

Author: mtvec

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -170,6 +170,7 @@ class MCPlusBuilder {
   /// MCPlusBuilder classes must use convert/lower/create* interfaces instead.
   void setTailCall(MCInst &Inst);
 
+public:
   /// Transfer annotations from \p SrcInst to \p DstInst.
   void moveAnnotations(MCInst &&SrcInst, MCInst &DstInst) const {
     assert(!getAnnotationInst(DstInst) &&
@@ -182,7 +183,6 @@ class MCPlusBuilder {
     removeAnnotationInst(SrcInst);
   }
 
-public:
   /// Return iterator range covering def operands.
   iterator_range<MCInst::iterator> defOperands(MCInst &Inst) const {
     return make_range(Inst.begin(),

bolt/lib/Passes/FixRISCVCallsPass.cpp

@@ -19,13 +19,15 @@ void FixRISCVCallsPass::runOnFunction(BinaryFunction &BF) {
         auto *Target = MIB->getTargetSymbol(*II);
         assert(Target && "Cannot find call target");
 
+        MCInst OldCall = *II;
         auto L = BC.scopeLock();
 
         if (MIB->isTailCall(*II))
           MIB->createTailCall(*II, Target, Ctx);
         else
           MIB->createCall(*II, Target, Ctx);
 
+        MIB->moveAnnotations(std::move(OldCall), *II);
         ++II;
         continue;
       }
@@ -39,8 +41,19 @@ void FixRISCVCallsPass::runOnFunction(BinaryFunction &BF) {
         auto *Target = MIB->getTargetSymbol(*II);
         assert(Target && "Cannot find call target");
 
+        MCInst OldCall = *NextII;
         auto L = BC.scopeLock();
         MIB->createCall(*II, Target, Ctx);
+        MIB->moveAnnotations(std::move(OldCall), *II);
+
+        // The original offset was set on the jalr of the auipc+jalr pair. Since
+        // the whole pair is replaced by a call, adjust the offset by -4 (the
+        // size of a auipc).
+        if (std::optional<uint32_t> Offset = MIB->getOffset(*II)) {
+          assert(*Offset >= 4 && "Illegal jalr offset");
+          MIB->setOffset(*II, *Offset - 4);
+        }
+
         II = BB.eraseInstruction(NextII);
         continue;
       }

bolt/test/RISCV/call-annotations.s

@@ -0,0 +1,33 @@
+/// Test that annotations are properly carried over to fixed calls.
+/// Note that --enable-bat is used to force offsets to be kept.
+
+// RUN: llvm-mc -triple riscv64 -filetype obj -o %t.o %s
+// RUN: ld.lld --emit-relocs -o %t %t.o
+// RUN: llvm-bolt --enable-bat --print-cfg --print-fix-riscv-calls \
+// RUN:     --print-only=_start -o /dev/null %t | FileCheck %s
+
+  .text
+  .global f
+  .p2align 1
+f:
+  ret
+  .size f, .-f
+
+// CHECK-LABEL: Binary Function "_start" after building cfg {
+// CHECK:      auipc ra, f
+// CHECK-NEXT: jalr ra, -4(ra) # Offset: 4
+// CHECK-NEXT: jal ra, f # Offset: 8
+// CHECK-NEXT: jal zero, f # TAILCALL  # Offset: 12
+
+// CHECK-LABEL: Binary Function "_start" after fix-riscv-calls {
+// CHECK:      call f # Offset: 0
+// CHECK-NEXT: call f # Offset: 8
+// CHECK-NEXT: tail f # TAILCALL   # Offset: 12
+
+  .globl _start
+  .p2align 1
+_start:
+  call f
+  jal f
+  jal zero, f
+  .size _start, .-_start