[llvm-branch-commits] [BOLT][BAT] Fix translate for branches added by BOLT (PR #90811)

[Amir Ayupov via llvm-branch-commits]

================
@@ -478,18 +478,34 @@ uint64_t BoltAddressTranslation::translate(uint64_t FuncAddress,
     return Offset;
 
   const MapTy &Map = Iter->second;
+  if (IsBranchSrc) {
+    // Try exact lookup first
+    auto KeyVal = Map.find(Offset);
+    if (KeyVal != Map.end() && KeyVal->second & BRANCHENTRY)
+      return KeyVal->second >> 1;
+  }
   auto KeyVal = Map.upper_bound(Offset);
   if (KeyVal == Map.begin())
     return Offset;
 
   --KeyVal;
 
   const uint32_t Val = KeyVal->second >> 1; // dropping BRANCHENTRY bit
-  // Branch source addresses are translated to the first instruction of the
-  // source BB to avoid accounting for modifications BOLT may have made in the
-  // BB regarding deletion/addition of instructions.
-  if (IsBranchSrc)
-    return Val;
+  if (IsBranchSrc) {
----------------
aaupov wrote:

Agree we can do an upper_bound followed by a comparison.

https://github.com/llvm/llvm-project/pull/90811