[Profile] Enable binary profile correlation for Mach-O binaries (#179937)

The existing ELF/COFF correlation code mostly "just works" on Mach-O
files, with one gotcha: on disk, the pointers in `__llvm_covdata` are
stored in an encoded format due to dyld fixup chains. (In memory, they
would normally be fixed up at load time in a running binary, but the
correlator only looks at the on-disk values.)

LLVM's Mach-O reader knows how to decode the format, so this patch walks
the fixup table to create a set of mappings that the correlator can use
to resolve the values.

rdar://168259786

compiler-rt/test/profile/instrprof-binary-correlate.c

@@ -1,4 +1,4 @@
-// REQUIRES: linux || windows
+// REQUIRES: linux || windows || darwin
 // Default
 // RUN: %clang -o %t.normal -fprofile-instr-generate -fcoverage-mapping %S/Inputs/instrprof-debug-info-correlate-main.cpp %S/Inputs/instrprof-debug-info-correlate-foo.cpp
 // RUN: env LLVM_PROFILE_FILE=%t.profraw %run %t.normal

llvm/include/llvm/ProfileData/InstrProfCorrelator.h

@@ -12,6 +12,7 @@
 #ifndef LLVM_PROFILEDATA_INSTRPROFCORRELATOR_H
 #define LLVM_PROFILEDATA_INSTRPROFCORRELATOR_H
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
 #include "llvm/Debuginfod/BuildIDFetcher.h"
@@ -77,7 +78,7 @@ public:
 protected:
   struct Context {
     LLVM_ABI static llvm::Expected<std::unique_ptr<Context>>
-    get(std::unique_ptr<MemoryBuffer> Buffer, const object::ObjectFile &Obj,
+    get(std::unique_ptr<MemoryBuffer> Buffer, object::ObjectFile &Obj,
         ProfCorrelatorKind FileKind);
     std::unique_ptr<MemoryBuffer> Buffer;
     /// The address range of the __llvm_prf_cnts section.
@@ -89,6 +90,10 @@ protected:
     const char *DataEnd;
     const char *NameStart;
     size_t NameSize;
+    /// Resolved values for Mach-O linker fixup chains when FileKind is Binary.
+    /// The mapping is from an address relative to the start of __llvm_covdata,
+    /// to the resolved pointer value at that address.
+    llvm::DenseMap<uint64_t, uint64_t> MachOFixups;
     /// True if target and host have different endian orders.
     bool ShouldSwapBytes;
   };

llvm/lib/ProfileData/InstrProfCorrelator.cpp

@@ -54,12 +54,13 @@ const char *InstrProfCorrelator::NumCountersAttributeName = "Num Counters";
 
 llvm::Expected<std::unique_ptr<InstrProfCorrelator::Context>>
 InstrProfCorrelator::Context::get(std::unique_ptr<MemoryBuffer> Buffer,
-                                  const object::ObjectFile &Obj,
+                                  object::ObjectFile &Obj,
                                   ProfCorrelatorKind FileKind) {
   auto C = std::make_unique<Context>();
   auto CountersSection = getInstrProfSection(Obj, IPSK_cnts);
   if (auto Err = CountersSection.takeError())
     return std::move(Err);
+  Triple::ObjectFormatType ObjFormat = Obj.getTripleObjectFormat();
   if (FileKind == InstrProfCorrelator::BINARY) {
     auto DataSection = getInstrProfSection(Obj, IPSK_covdata);
     if (auto Err = DataSection.takeError())
@@ -77,13 +78,32 @@ InstrProfCorrelator::Context::get(std::unique_ptr<MemoryBuffer> Buffer,
     C->DataEnd = DataOrErr->data() + DataOrErr->size();
     C->NameStart = NameOrErr->data();
     C->NameSize = NameOrErr->size();
+
+    if (ObjFormat == Triple::MachO) {
+      std::string FullSectionName =
+          getInstrProfSectionName(IPSK_covdata, ObjFormat);
+      SmallVector<StringRef, 3> SegmentAndSection;
+      StringRef(FullSectionName).split(SegmentAndSection, ',', 2);
+      auto *MachO = static_cast<object::MachOObjectFile *>(&Obj);
+      Error Err = Error::success();
+      for (const object::MachOChainedFixupEntry &Entry :
+           MachO->fixupTable(Err)) {
+        if (Entry.isRebase() && Entry.segmentName() == SegmentAndSection[0] &&
+            Entry.sectionName() == SegmentAndSection[1]) {
+          C->MachOFixups[Entry.address() - DataSection->getAddress()] =
+              Entry.pointerValue();
+        }
+      }
+      if (Err)
+        return std::move(Err);
+    }
   }
   C->Buffer = std::move(Buffer);
   C->CountersSectionStart = CountersSection->getAddress();
   C->CountersSectionEnd = C->CountersSectionStart + CountersSection->getSize();
   // In COFF object file, there's a null byte at the beginning of the counter
   // section which doesn't exist in raw profile.
-  if (Obj.getTripleObjectFormat() == Triple::COFF)
+  if (ObjFormat == Triple::COFF)
     ++C->CountersSectionStart;
 
   C->ShouldSwapBytes = Obj.isLittleEndian() != sys::IsLittleEndianHost;
@@ -218,11 +238,11 @@ InstrProfCorrelatorImpl<IntPtrT>::get(
         instrprof_error::unable_to_correlate_profile,
         "unsupported debug info format (only DWARF is supported)");
   }
-  if (Obj.isELF() || Obj.isCOFF())
+  if (Obj.isELF() || Obj.isCOFF() || Obj.isMachO())
     return std::make_unique<BinaryInstrProfCorrelator<IntPtrT>>(std::move(Ctx));
   return make_error<InstrProfError>(
       instrprof_error::unable_to_correlate_profile,
-      "unsupported binary format (only ELF and COFF are supported)");
+      "unsupported binary format (only ELF, COFF, and Mach-O are supported)");
 }
 
 template <class IntPtrT>
@@ -478,6 +498,16 @@ void BinaryInstrProfCorrelator<IntPtrT>::correlateProfileDataImpl(
     uint64_t CounterPtr = this->template maybeSwap<IntPtrT>(I->CounterPtr);
     uint64_t CountersStart = this->Ctx->CountersSectionStart;
     uint64_t CountersEnd = this->Ctx->CountersSectionEnd;
+    if (!this->Ctx->MachOFixups.empty()) {
+      uint64_t Offset = (uint64_t)&I->CounterPtr - (uint64_t)DataStart;
+      auto It = this->Ctx->MachOFixups.find(Offset);
+      if (It != this->Ctx->MachOFixups.end()) {
+        CounterPtr = It->second;
+      } else if (UnlimitedWarnings || ++NumSuppressedWarnings < 1) {
+        WithColor::warning() << format(
+            "Mach-O fixup not found for covdata offset 0x%llx\n", Offset);
+      }
+    }
     if (CounterPtr < CountersStart || CounterPtr >= CountersEnd) {
       if (UnlimitedWarnings || ++NumSuppressedWarnings < 1) {
         WithColor::warning()