From 54b17fa4eefb7ccb00fba8189e92fac2d46cc481 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov@fb.com>
Date: Sun, 12 May 2024 14:26:41 -0700
Subject: [PATCH] [BOLT] Preserve Offset annotation in fixDoubleJumps (#91898)

Offset annotation was missed when optimizing an unconditional branch to
a tail call.

Test Plan: update bb-with-two-tail-calls.s
---
 bolt/lib/Passes/BinaryPasses.cpp       | 3 +++
 bolt/test/X86/bb-with-two-tail-calls.s | 6 +++---
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index df6dbcddeed5..867f977cebca 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -715,6 +715,9 @@ static uint64_t fixDoubleJumps(BinaryFunction &Function, bool MarkInvalid) {
           Pred->removeSuccessor(&BB);
           Pred->eraseInstruction(Pred->findInstruction(Branch));
           Pred->addTailCallInstruction(SuccSym);
+          MCInst *TailCall = Pred->getLastNonPseudoInstr();
+          assert(TailCall);
+          MIB->setOffset(*TailCall, BB.getOffset());
         } else {
           return false;
         }
diff --git a/bolt/test/X86/bb-with-two-tail-calls.s b/bolt/test/X86/bb-with-two-tail-calls.s
index caad7b3d735f..bb2b0cd4cc23 100644
--- a/bolt/test/X86/bb-with-two-tail-calls.s
+++ b/bolt/test/X86/bb-with-two-tail-calls.s
@@ -9,11 +9,11 @@
 # RUN: llvm-strip --strip-unneeded %t.o
 # RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -nostdlib
 # RUN: llvm-bolt %t.exe -o %t.out --data %t.fdata --lite=0 --dyno-stats \
-# RUN:    --print-sctc --print-only=_start 2>&1 | FileCheck %s
+# RUN:    --print-sctc --print-only=_start -enable-bat 2>&1 | FileCheck %s
 # CHECK-NOT: Assertion `BranchInfo.size() == 2 && "could only be called for blocks with 2 successors"' failed.
 # Two tail calls in the same basic block after SCTC:
-# CHECK:         {{.*}}:   ja      {{.*}} # TAILCALL  # CTCTakenCount: {{.*}}
-# CHECK-NEXT:    {{.*}}:   jmp     {{.*}} # TAILCALL
+# CHECK:         {{.*}}:   ja      {{.*}} # TAILCALL # Offset: 7 # CTCTakenCount: 4
+# CHECK-NEXT:    {{.*}}:   jmp     {{.*}} # TAILCALL # Offset: 12
 
   .globl _start
 _start: