[llvm-profdata] Fix binary ids with multiple raw profiles in a single… (#72740)

Save binary ids when iterating through `RawInstrProfReader`.

Fixes #72699.

Author: ZequanWu

compiler-rt/test/profile/Linux/binary-id.c

@@ -17,16 +17,35 @@
 // RUN: llvm-profdata show --binary-ids  %t.profdir/default_*.profraw > %t.profraw.out
 // RUN: FileCheck %s --check-prefix=BINARY-ID-MERGE-PROF < %t.profraw.out
 
-// RUN: llvm-profdata merge -o %t.profdata %t.profraw %t.profraw
+// RUN: llvm-profdata merge -o %t.profdata %t.profdir/default_*.profraw
 // RUN: llvm-profdata show --binary-ids %t.profdata > %t.profdata.out
-// RUN: FileCheck %s --check-prefix=BINARY-ID-INDEXED-PROF < %t.profraw.out
+// RUN: FileCheck %s --check-prefix=BINARY-ID-INDEXED-PROF < %t.profdata.out
 
+// Test raw profiles with shared libraries.
+// RUN: split-file %s %t.dir
+// RUN: %clang_profgen -Wl,--build-id -fpic -shared -O2 %t.dir/foo.c -o %t.dir/libfoo.so
+// RUN: %clang_profgen -Wl,--build-id -fpic -shared -O2 %t.dir/bar.c -o %t.dir/libbar.so
+// RUN: %clang_profgen -Wl,--build-id -O2 %t.dir/main.c %t.dir/libfoo.so %t.dir/libbar.so -o %t
+// RUN: env LLVM_PROFILE_FILE=%t.profraw LD_LIBRARY_PATH=%t.dir %run %t
+// RUN: llvm-profdata show --binary-ids %t.profraw > %t.profraw.out
+// RUN: llvm-profdata merge -o %t.profdata %t.profraw
+// RUN: FileCheck %s --check-prefix=BINARY-ID-SHARE-RAW-PROF < %t.profraw.out
+
+// RUN: llvm-profdata merge -o %t.profdata %t.profraw
+// RUN: llvm-profdata show --binary-ids %t.profdata > %t.profdata.out
+// RUN: FileCheck %s --check-prefix=BINARY-ID-SHARE-INDEXED-PROF < %t.profraw.out
+
+//--- foo.c
 void foo() {
 }
 
+//--- bar.c
 void bar() {
 }
 
+//--- main.c
+void foo();
+void bar();
 int main() {
   foo();
   bar();
@@ -59,3 +78,21 @@ int main() {
 // BINARY-ID-INDEXED-PROF-NEXT: Maximum internal block count: 0
 // BINARY-ID-INDEXED-PROF-NEXT: Binary IDs:
 // BINARY-ID-INDEXED-PROF-NEXT: {{[0-9a-f]+}}
+
+// BINARY-ID-SHARE-RAW-PROF: Instrumentation level: Front-end
+// BINARY-ID-SHARE-RAW-PROF-NEXT: Total functions: 3
+// BINARY-ID-SHARE-RAW-PROF-NEXT: Maximum function count: 1
+// BINARY-ID-SHARE-RAW-PROF-NEXT: Maximum internal block count: 0
+// BINARY-ID-SHARE-RAW-PROF-NEXT: Binary IDs:
+// BINARY-ID-SHARE-RAW-PROF-NEXT: {{[0-9a-f]+}}
+// BINARY-ID-SHARE-RAW-PROF-NEXT: {{[0-9a-f]+}}
+// BINARY-ID-SHARE-RAW-PROF-NEXT: {{[0-9a-f]+}}
+
+// BINARY-ID-SHARE-INDEXED-PROF: Instrumentation level: Front-end
+// BINARY-ID-SHARE-INDEXED-PROF-NEXT: Total functions: 3
+// BINARY-ID-SHARE-INDEXED-PROF-NEXT: Maximum function count: 1
+// BINARY-ID-SHARE-INDEXED-PROF-NEXT: Maximum internal block count: 0
+// BINARY-ID-SHARE-INDEXED-PROF-NEXT: Binary IDs:
+// BINARY-ID-SHARE-INDEXED-PROF-NEXT: {{[0-9a-f]+}}
+// BINARY-ID-SHARE-INDEXED-PROF-NEXT: {{[0-9a-f]+}}
+// BINARY-ID-SHARE-INDEXED-PROF-NEXT: {{[0-9a-f]+}}

llvm/include/llvm/ProfileData/InstrProfReader.h

@@ -340,11 +340,7 @@ class RawInstrProfReader : public InstrProfReader {
   const uint8_t *ValueDataStart;
   uint32_t ValueKindLast;
   uint32_t CurValueDataSize;
-
-  /// Total size of binary ids.
-  uint64_t BinaryIdsSize{0};
-  /// Start address of binary id length and data pairs.
-  const uint8_t *BinaryIdsStart;
+  std::vector<llvm::object::BuildID> BinaryIds;
 
   std::function<void(Error)> Warn;
 

llvm/lib/ProfileData/InstrProfReader.cpp

@@ -142,27 +142,15 @@ readBinaryIdsInternal(const MemoryBuffer &DataBuffer,
   return Error::success();
 }
 
-static Error printBinaryIdsInternal(raw_ostream &OS,
-                                    const MemoryBuffer &DataBuffer,
-                                    uint64_t BinaryIdsSize,
-                                    const uint8_t *BinaryIdsStart,
-                                    llvm::endianness Endian) {
-  if (BinaryIdsSize == 0)
-    return Error::success();
-
-  std::vector<llvm::object::BuildID> BinaryIds;
-  if (Error E = readBinaryIdsInternal(DataBuffer, BinaryIdsSize, BinaryIdsStart,
-                                      BinaryIds, Endian))
-    return E;
-
+static void
+printBinaryIdsInternal(raw_ostream &OS,
+                       std::vector<llvm::object::BuildID> &BinaryIds) {
   OS << "Binary IDs: \n";
   for (auto BI : BinaryIds) {
     for (uint64_t I = 0; I < BI.size(); I++)
       OS << format("%02x", BI[I]);
     OS << "\n";
   }
-
-  return Error::success();
 }
 
 Expected<std::unique_ptr<InstrProfReader>>
@@ -573,9 +561,20 @@ Error RawInstrProfReader<IntPtrT>::readHeader(
   if (!useCorrelate() && Correlator)
     return error(instrprof_error::unexpected_debug_info_for_correlation);
 
-  BinaryIdsSize = swap(Header.BinaryIdsSize);
-  if (BinaryIdsSize % sizeof(uint64_t))
+  uint64_t BinaryIdSize = swap(Header.BinaryIdsSize);
+  // Binary id start just after the header if exists.
+  const uint8_t *BinaryIdStart =
+      reinterpret_cast<const uint8_t *>(&Header) + sizeof(RawInstrProf::Header);
+  const uint8_t *BinaryIdEnd = BinaryIdStart + BinaryIdSize;
+  const uint8_t *BufferEnd = (const uint8_t *)DataBuffer->getBufferEnd();
+  if (BinaryIdSize % sizeof(uint64_t) || BinaryIdEnd > BufferEnd)
     return error(instrprof_error::bad_header);
+  if (BinaryIdSize != 0) {
+    if (Error Err =
+            readBinaryIdsInternal(*DataBuffer, BinaryIdSize, BinaryIdStart,
+                                  BinaryIds, getDataEndianness()))
+      return Err;
+  }
 
   CountersDelta = swap(Header.CountersDelta);
   BitmapDelta = swap(Header.BitmapDelta);
@@ -593,7 +592,7 @@ Error RawInstrProfReader<IntPtrT>::readHeader(
   auto PaddingSize = getNumPaddingBytes(NamesSize);
 
   // Profile data starts after profile header and binary ids if exist.
-  ptrdiff_t DataOffset = sizeof(RawInstrProf::Header) + BinaryIdsSize;
+  ptrdiff_t DataOffset = sizeof(RawInstrProf::Header) + BinaryIdSize;
   ptrdiff_t CountersOffset = DataOffset + DataSize + PaddingBytesBeforeCounters;
   ptrdiff_t BitmapOffset =
       CountersOffset + CountersSize + PaddingBytesAfterCounters;
@@ -622,19 +621,12 @@ Error RawInstrProfReader<IntPtrT>::readHeader(
     NamesEnd = NamesStart + NamesSize;
   }
 
-  // Binary ids start just after the header.
-  BinaryIdsStart =
-      reinterpret_cast<const uint8_t *>(&Header) + sizeof(RawInstrProf::Header);
   CountersStart = Start + CountersOffset;
   CountersEnd = CountersStart + CountersSize;
   BitmapStart = Start + BitmapOffset;
   BitmapEnd = BitmapStart + NumBitmapBytes;
   ValueDataStart = reinterpret_cast<const uint8_t *>(Start + ValueDataOffset);
 
-  const uint8_t *BufferEnd = (const uint8_t *)DataBuffer->getBufferEnd();
-  if (BinaryIdsStart + BinaryIdsSize > BufferEnd)
-    return error(instrprof_error::bad_header);
-
   std::unique_ptr<InstrProfSymtab> NewSymtab = std::make_unique<InstrProfSymtab>();
   if (Error E = createSymtab(*NewSymtab))
     return E;
@@ -832,14 +824,16 @@ Error RawInstrProfReader<IntPtrT>::readNextRecord(NamedInstrProfRecord &Record)
 template <class IntPtrT>
 Error RawInstrProfReader<IntPtrT>::readBinaryIds(
     std::vector<llvm::object::BuildID> &BinaryIds) {
-  return readBinaryIdsInternal(*DataBuffer, BinaryIdsSize, BinaryIdsStart,
-                               BinaryIds, getDataEndianness());
+  BinaryIds.insert(BinaryIds.begin(), this->BinaryIds.begin(),
+                   this->BinaryIds.end());
+  return Error::success();
 }
 
 template <class IntPtrT>
 Error RawInstrProfReader<IntPtrT>::printBinaryIds(raw_ostream &OS) {
-  return printBinaryIdsInternal(OS, *DataBuffer, BinaryIdsSize, BinaryIdsStart,
-                                getDataEndianness());
+  if (!BinaryIds.empty())
+    printBinaryIdsInternal(OS, BinaryIds);
+  return Error::success();
 }
 
 namespace llvm {
@@ -1475,8 +1469,11 @@ Error IndexedInstrProfReader::readBinaryIds(
 }
 
 Error IndexedInstrProfReader::printBinaryIds(raw_ostream &OS) {
-  return printBinaryIdsInternal(OS, *DataBuffer, BinaryIdsSize, BinaryIdsStart,
-                                llvm::endianness::little);
+  std::vector<llvm::object::BuildID> BinaryIds;
+  if (Error E = readBinaryIds(BinaryIds))
+    return E;
+  printBinaryIdsInternal(OS, BinaryIds);
+  return Error::success();
 }
 
 void InstrProfReader::accumulateCounts(CountSumOrPercent &Sum, bool IsCS) {