[llvm-branch-commits] [llvm] [BOLT] Map branch source address to the containing basic block in BAT YAML (PR #91273)

[Maksim Panchenko via llvm-branch-commits]

================
@@ -2378,21 +2378,24 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
         return CSI;
       };
 
+      // Lookup containing basic block offset and index
+      auto getBlock = [&BlockMap](uint32_t Offset) {
+        auto BlockIt = BlockMap.upper_bound(Offset);
+        assert(BlockIt != BlockMap.begin());
+        --BlockIt;
+        return std::pair(BlockIt->first, BlockIt->second.getBBIndex());
+      };
+
       for (const auto &[FromOffset, SuccKV] : Branches.IntraIndex) {
-        if (!BlockMap.isInputBlock(FromOffset))
-          continue;
-        const unsigned Index = BlockMap.getBBIndex(FromOffset);
+        const auto &[_, Index] = getBlock(FromOffset);
----------------
maksfb wrote:

Even though we generate BAT in BOLT, when we view the invocation of BOLT on a binary with embedded BAT, such input should be considered an external and potentially malformed data. In this case, assertions will not provide adequate enough protection since we can build BOLT without them.

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