1 files changed, 69 insertions, 5 deletions

diff --git a/clang/lib/CodeGen/CoverageMappingGen.cpp b/clang/lib/CodeGen/CoverageMappingGen.cpp
index 78b268f..6729c7f 100644
--- a/clang/lib/CodeGen/CoverageMappingGen.cpp
+++ b/clang/lib/CodeGen/CoverageMappingGen.cpp
@@ -35,8 +35,40 @@ using namespace clang;
 using namespace CodeGen;
 using namespace llvm::coverage;
 
+CoverageSourceInfo *
+CoverageMappingModuleGen::setUpCoverageCallbacks(Preprocessor &PP) {
+  CoverageSourceInfo *CoverageInfo = new CoverageSourceInfo;
+  PP.addPPCallbacks(std::unique_ptr<PPCallbacks>(CoverageInfo));
+  PP.addCommentHandler(CoverageInfo);
+  PP.setPreprocessToken(true);
+  PP.setTokenWatcher([CoverageInfo](clang::Token Tok) {
+    // Update previous token location.
+    CoverageInfo->PrevTokLoc = Tok.getLocation();
+    CoverageInfo->updateNextTokLoc(Tok.getLocation());
+  });
+  return CoverageInfo;
+}
+
 void CoverageSourceInfo::SourceRangeSkipped(SourceRange Range, SourceLocation) {
-  SkippedRanges.push_back(Range);
+  SkippedRanges.push_back({Range});
+}
+
+bool CoverageSourceInfo::HandleComment(Preprocessor &PP, SourceRange Range) {
+  if (PrevTokLoc.isValid() && PrevTokLoc == BeforeCommentLoc)
+    SkippedRanges.back().Range.setEnd(Range.getEnd());
+  else {
+    SkippedRanges.push_back({Range, PrevTokLoc});
+    BeforeCommentLoc = PrevTokLoc;
+  }
+  LastCommentIndex = SkippedRanges.size() - 1;
+  return false;
+}
+
+void CoverageSourceInfo::updateNextTokLoc(SourceLocation Loc) {
+  if (LastCommentIndex) {
+    SkippedRanges[LastCommentIndex.getValue()].NextTokLoc = Loc;
+    LastCommentIndex = None;
+  }
 }
 
 namespace {
@@ -274,8 +306,34 @@ public:
     return None;
   }
 
+  /// This shrinks the skipped range if it spans a line that contains a
+  /// non-comment token. If shrinking the skipped range would make it empty,
+  /// this returns None.
+  Optional<SpellingRegion> adjustSkippedRange(SourceManager &SM,
+                                              SpellingRegion SR,
+                                              SourceLocation PrevTokLoc,
+                                              SourceLocation NextTokLoc) {
+    // If Range begin location is invalid, it's not a comment region.
+    if (PrevTokLoc.isInvalid())
+      return SR;
+    unsigned PrevTokLine = SM.getSpellingLineNumber(PrevTokLoc);
+    unsigned NextTokLine = SM.getSpellingLineNumber(NextTokLoc);
+    SpellingRegion newSR(SR);
+    if (SR.LineStart == PrevTokLine) {
+      newSR.LineStart = SR.LineStart + 1;
+      newSR.ColumnStart = 1;
+    }
+    if (SR.LineEnd == NextTokLine) {
+      newSR.LineEnd = SR.LineEnd - 1;
+      newSR.ColumnEnd = 1;
+    }
+    if (newSR.isInSourceOrder())
+      return newSR;
+    return None;
+  }
+
   /// Gather all the regions that were skipped by the preprocessor
-  /// using the constructs like #if.
+  /// using the constructs like #if or comments.
   void gatherSkippedRegions() {
     /// An array of the minimum lineStarts and the maximum lineEnds
     /// for mapping regions from the appropriate source files.
@@ -291,9 +349,10 @@ public:
     }
 
     auto SkippedRanges = CVM.getSourceInfo().getSkippedRanges();
-    for (const auto &I : SkippedRanges) {
-      auto LocStart = I.getBegin();
-      auto LocEnd = I.getEnd();
+    for (auto &I : SkippedRanges) {
+      SourceRange Range = I.Range;
+      auto LocStart = Range.getBegin();
+      auto LocEnd = Range.getEnd();
       assert(SM.isWrittenInSameFile(LocStart, LocEnd) &&
              "region spans multiple files");
 
@@ -301,6 +360,11 @@ public:
       if (!CovFileID)
         continue;
       SpellingRegion SR{SM, LocStart, LocEnd};
+      if (Optional<SpellingRegion> res =
+              adjustSkippedRange(SM, SR, I.PrevTokLoc, I.NextTokLoc))
+        SR = res.getValue();
+      else
+        continue;
       auto Region = CounterMappingRegion::makeSkipped(
           *CovFileID, SR.LineStart, SR.ColumnStart, SR.LineEnd, SR.ColumnEnd);
       // Make sure that we only collect the regions that are inside