1 files changed, 57 insertions, 18 deletions

diff --git a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
index a8cc308b4e3a..94c2bee3590c 100644
--- a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
+++ b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
@@ -186,6 +186,22 @@ Expected<int64_t> CounterMappingContext::evaluate(const Counter &C) const {
   llvm_unreachable("Unhandled CounterKind");
 }
 
+unsigned CounterMappingContext::getMaxCounterID(const Counter &C) const {
+  switch (C.getKind()) {
+  case Counter::Zero:
+    return 0;
+  case Counter::CounterValueReference:
+    return C.getCounterID();
+  case Counter::Expression: {
+    if (C.getExpressionID() >= Expressions.size())
+      return 0;
+    const auto &E = Expressions[C.getExpressionID()];
+    return std::max(getMaxCounterID(E.LHS), getMaxCounterID(E.RHS));
+  }
+  }
+  llvm_unreachable("Unhandled CounterKind");
+}
+
 void FunctionRecordIterator::skipOtherFiles() {
   while (Current != Records.end() && !Filename.empty() &&
          Filename != Current->Filenames[0])
@@ -203,6 +219,15 @@ ArrayRef<unsigned> CoverageMapping::getImpreciseRecordIndicesForFilename(
   return RecordIt->second;
 }
 
+static unsigned getMaxCounterID(const CounterMappingContext &Ctx,
+                                const CoverageMappingRecord &Record) {
+  unsigned MaxCounterID = 0;
+  for (const auto &Region : Record.MappingRegions) {
+    MaxCounterID = std::max(MaxCounterID, Ctx.getMaxCounterID(Region.Count));
+  }
+  return MaxCounterID;
+}
+
 Error CoverageMapping::loadFunctionRecord(
     const CoverageMappingRecord &Record,
     IndexedInstrProfReader &ProfileReader) {
@@ -227,7 +252,7 @@ Error CoverageMapping::loadFunctionRecord(
       return Error::success();
     } else if (IPE != instrprof_error::unknown_function)
       return make_error<InstrProfError>(IPE);
-    Counts.assign(Record.MappingRegions.size(), 0);
+    Counts.assign(getMaxCounterID(Ctx, Record) + 1, 0);
   }
   Ctx.setCounts(Counts);
 
@@ -281,21 +306,29 @@ Error CoverageMapping::loadFunctionRecord(
   return Error::success();
 }
 
-Expected<std::unique_ptr<CoverageMapping>> CoverageMapping::load(
+// This function is for memory optimization by shortening the lifetimes
+// of CoverageMappingReader instances.
+Error CoverageMapping::loadFromReaders(
     ArrayRef<std::unique_ptr<CoverageMappingReader>> CoverageReaders,
-    IndexedInstrProfReader &ProfileReader) {
-  auto Coverage = std::unique_ptr<CoverageMapping>(new CoverageMapping());
-
+    IndexedInstrProfReader &ProfileReader, CoverageMapping &Coverage) {
   for (const auto &CoverageReader : CoverageReaders) {
     for (auto RecordOrErr : *CoverageReader) {
       if (Error E = RecordOrErr.takeError())
-        return std::move(E);
+        return E;
       const auto &Record = *RecordOrErr;
-      if (Error E = Coverage->loadFunctionRecord(Record, ProfileReader))
-        return std::move(E);
+      if (Error E = Coverage.loadFunctionRecord(Record, ProfileReader))
+        return E;
     }
   }
+  return Error::success();
+}
 
+Expected<std::unique_ptr<CoverageMapping>> CoverageMapping::load(
+    ArrayRef<std::unique_ptr<CoverageMappingReader>> CoverageReaders,
+    IndexedInstrProfReader &ProfileReader) {
+  auto Coverage = std::unique_ptr<CoverageMapping>(new CoverageMapping());
+  if (Error E = loadFromReaders(CoverageReaders, ProfileReader, *Coverage))
+    return std::move(E);
   return std::move(Coverage);
 }
 
@@ -311,23 +344,26 @@ static Error handleMaybeNoDataFoundError(Error E) {
 
 Expected<std::unique_ptr<CoverageMapping>>
 CoverageMapping::load(ArrayRef<StringRef> ObjectFilenames,
-                      StringRef ProfileFilename, ArrayRef<StringRef> Arches) {
+                      StringRef ProfileFilename, ArrayRef<StringRef> Arches,
+                      StringRef CompilationDir) {
   auto ProfileReaderOrErr = IndexedInstrProfReader::create(ProfileFilename);
   if (Error E = ProfileReaderOrErr.takeError())
     return std::move(E);
   auto ProfileReader = std::move(ProfileReaderOrErr.get());
+  auto Coverage = std::unique_ptr<CoverageMapping>(new CoverageMapping());
+  bool DataFound = false;
 
-  SmallVector<std::unique_ptr<CoverageMappingReader>, 4> Readers;
-  SmallVector<std::unique_ptr<MemoryBuffer>, 4> Buffers;
   for (const auto &File : llvm::enumerate(ObjectFilenames)) {
-    auto CovMappingBufOrErr = MemoryBuffer::getFileOrSTDIN(File.value());
+    auto CovMappingBufOrErr = MemoryBuffer::getFileOrSTDIN(
+        File.value(), /*IsText=*/false, /*RequiresNullTerminator=*/false);
     if (std::error_code EC = CovMappingBufOrErr.getError())
       return errorCodeToError(EC);
     StringRef Arch = Arches.empty() ? StringRef() : Arches[File.index()];
     MemoryBufferRef CovMappingBufRef =
         CovMappingBufOrErr.get()->getMemBufferRef();
-    auto CoverageReadersOrErr =
-        BinaryCoverageReader::create(CovMappingBufRef, Arch, Buffers);
+    SmallVector<std::unique_ptr<MemoryBuffer>, 4> Buffers;
+    auto CoverageReadersOrErr = BinaryCoverageReader::create(
+        CovMappingBufRef, Arch, Buffers, CompilationDir);
     if (Error E = CoverageReadersOrErr.takeError()) {
       E = handleMaybeNoDataFoundError(std::move(E));
       if (E)
@@ -335,15 +371,19 @@ CoverageMapping::load(ArrayRef<StringRef> ObjectFilenames,
       // E == success (originally a no_data_found error).
       continue;
     }
+
+    SmallVector<std::unique_ptr<CoverageMappingReader>, 4> Readers;
     for (auto &Reader : CoverageReadersOrErr.get())
       Readers.push_back(std::move(Reader));
-    Buffers.push_back(std::move(CovMappingBufOrErr.get()));
+    DataFound |= !Readers.empty();
+    if (Error E = loadFromReaders(Readers, *ProfileReader, *Coverage))
+      return std::move(E);
   }
   // If no readers were created, either no objects were provided or none of them
   // had coverage data. Return an error in the latter case.
-  if (Readers.empty() && !ObjectFilenames.empty())
+  if (!DataFound && !ObjectFilenames.empty())
     return make_error<CoverageMapError>(coveragemap_error::no_data_found);
-  return load(Readers, *ProfileReader);
+  return std::move(Coverage);
 }
 
 namespace {
@@ -794,7 +834,6 @@ LineCoverageStats::LineCoverageStats(
     ExecutionCount = WrappedSegment->Count;
   if (!MinRegionCount)
     return;
-  ExecutionCount = 0;
   for (const auto *LS : LineSegments)
     if (isStartOfRegion(LS))
       ExecutionCount = std::max(ExecutionCount, LS->Count);