[CoverageMapping] Handle gaps in counter IDs for source-based coverage

For source-based coverage, the frontend sets the counter IDs and the
constraints of counter IDs is not defined.  For e.g., the Rust frontend
until recently had a reserved counter #0
(https://github.com/rust-lang/rust/pull/83774).  Rust coverage
instrumentation also creates counters on edges in addition to basic
blocks.  Some functions may have more counters than regions.

This breaks an assumption in CoverageMapping.cpp where the number of
counters in a function is assumed to be bounded by the number of
regions:
  Counts.assign(Record.MappingRegions.size(), 0);

This assumption causes CounterMappingContext::evaluate() to fail since
there are not enough counter values created in the above call to
`Counts.assign`.  Consequently, some uncovered functions are not
reported in coverage reports.

This change walks a Function's CoverageMappingRecord to find the maximum
counter ID, and uses it to initialize the counter array when instrprof
records are missing for a function in sparse profiles.

Differential Revision: https://reviews.llvm.org/D101780

include/llvm/ProfileData/Coverage/CoverageMapping.h

@@ -334,6 +334,8 @@ public:
   /// Return the number of times that a region of code associated with this
   /// counter was executed.
   Expected<int64_t> evaluate(const Counter &C) const;
+
+  unsigned getMaxCounterID(const Counter &C) const;
 };
 
 /// Code coverage information for a single function.

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);
 

unittests/ProfileData/CoverageMappingTest.cpp

@@ -62,6 +62,7 @@ struct OutputFunctionCoverageData {
   uint64_t Hash;
   std::vector<StringRef> Filenames;
   std::vector<CounterMappingRegion> Regions;
+  std::vector<CounterExpression> Expressions;
 
   OutputFunctionCoverageData() : Hash(0) {}
 
@@ -78,7 +79,7 @@ struct OutputFunctionCoverageData {
     Record.FunctionName = Name;
     Record.FunctionHash = Hash;
     Record.Filenames = Filenames;
-    Record.Expressions = {};
+    Record.Expressions = Expressions;
     Record.MappingRegions = Regions;
   }
 };
@@ -111,6 +112,7 @@ struct InputFunctionCoverageData {
   std::string Name;
   uint64_t Hash;
   std::vector<CounterMappingRegion> Regions;
+  std::vector<CounterExpression> Expressions;
 
   InputFunctionCoverageData(std::string Name, uint64_t Hash)
       : Name(std::move(Name)), Hash(Hash) {}
@@ -189,13 +191,17 @@ struct CoverageMappingTest : ::testing::TestWithParam<std::tuple<bool, bool>> {
         LS, CS, LE, CE));
   }
 
+  void addExpression(CounterExpression CE) {
+    InputFunctions.back().Expressions.push_back(CE);
+  }
+
   std::string writeCoverageRegions(InputFunctionCoverageData &Data) {
     SmallVector<unsigned, 8> FileIDs(Data.ReverseVirtualFileMapping.size());
     for (const auto &E : Data.ReverseVirtualFileMapping)
       FileIDs[E.second] = E.first;
     std::string Coverage;
     llvm::raw_string_ostream OS(Coverage);
-    CoverageMappingWriter(FileIDs, None, Data.Regions).write(OS);
+    CoverageMappingWriter(FileIDs, Data.Expressions, Data.Regions).write(OS);
     return OS.str();
   }
 
@@ -207,10 +213,9 @@ struct CoverageMappingTest : ::testing::TestWithParam<std::tuple<bool, bool>> {
     Filenames.resize(Files.size() + 1);
     for (const auto &E : Files)
       Filenames[E.getValue()] = E.getKey().str();
-    std::vector<CounterExpression> Expressions;
     ArrayRef<std::string> FilenameRefs = llvm::makeArrayRef(Filenames);
     RawCoverageMappingReader Reader(Coverage, FilenameRefs, Data.Filenames,
-                                    Expressions, Data.Regions);
+                                    Data.Expressions, Data.Regions);
     EXPECT_THAT_ERROR(Reader.read(), Succeeded());
   }
 
@@ -796,6 +801,26 @@ TEST_P(CoverageMappingTest, combine_expansions) {
   EXPECT_EQ(CoverageSegment(5, 5, false), Segments[3]);
 }
 
+// Test that counters not associated with any code regions are allowed.
+TEST_P(CoverageMappingTest, non_code_region_counters) {
+  // No records in profdata
+
+  startFunction("func", 0x1234);
+  addCMR(Counter::getCounter(0), "file", 1, 1, 5, 5);
+  addCMR(Counter::getExpression(0), "file", 6, 1, 6, 5);
+  addExpression(CounterExpression(
+      CounterExpression::Add, Counter::getCounter(1), Counter::getCounter(2)));
+
+  EXPECT_THAT_ERROR(loadCoverageMapping(), Succeeded());
+
+  std::vector<std::string> Names;
+  for (const auto &Func : LoadedCoverage->getCoveredFunctions()) {
+    Names.push_back(Func.Name);
+    ASSERT_EQ(2U, Func.CountedRegions.size());
+  }
+  ASSERT_EQ(1U, Names.size());
+}
+
 TEST_P(CoverageMappingTest, strip_filename_prefix) {
   ProfileWriter.addRecord({"file1:func", 0x1234, {0}}, Err);