1 files changed, 145 insertions, 167 deletions

diff --git a/clang/lib/CodeGen/CoverageMappingGen.cpp b/clang/lib/CodeGen/CoverageMappingGen.cpp
index 5c25c204cc0b..8a11da600e4a 100644
--- a/clang/lib/CodeGen/CoverageMappingGen.cpp
+++ b/clang/lib/CodeGen/CoverageMappingGen.cpp
@@ -104,26 +104,21 @@ class SourceMappingRegion {
   /// The region's ending location.
   Optional<SourceLocation> LocEnd;
 
-  /// Whether this region should be emitted after its parent is emitted.
-  bool DeferRegion;
-
   /// Whether this region is a gap region. The count from a gap region is set
   /// as the line execution count if there are no other regions on the line.
   bool GapRegion;
 
 public:
   SourceMappingRegion(Counter Count, Optional<SourceLocation> LocStart,
-                      Optional<SourceLocation> LocEnd, bool DeferRegion = false,
-                      bool GapRegion = false)
-      : Count(Count), LocStart(LocStart), LocEnd(LocEnd),
-        DeferRegion(DeferRegion), GapRegion(GapRegion) {}
+                      Optional<SourceLocation> LocEnd, bool GapRegion = false)
+      : Count(Count), LocStart(LocStart), LocEnd(LocEnd), GapRegion(GapRegion) {
+  }
 
   SourceMappingRegion(Counter Count, Optional<Counter> FalseCount,
                       Optional<SourceLocation> LocStart,
-                      Optional<SourceLocation> LocEnd, bool DeferRegion = false,
-                      bool GapRegion = false)
+                      Optional<SourceLocation> LocEnd, bool GapRegion = false)
       : Count(Count), FalseCount(FalseCount), LocStart(LocStart),
-        LocEnd(LocEnd), DeferRegion(DeferRegion), GapRegion(GapRegion) {}
+        LocEnd(LocEnd), GapRegion(GapRegion) {}
 
   const Counter &getCounter() const { return Count; }
 
@@ -155,10 +150,6 @@ public:
     return *LocEnd;
   }
 
-  bool isDeferred() const { return DeferRegion; }
-
-  void setDeferred(bool Deferred) { DeferRegion = Deferred; }
-
   bool isGap() const { return GapRegion; }
 
   void setGap(bool Gap) { GapRegion = Gap; }
@@ -544,10 +535,6 @@ struct CounterCoverageMappingBuilder
   /// A stack of currently live regions.
   std::vector<SourceMappingRegion> RegionStack;
 
-  /// The currently deferred region: its end location and count can be set once
-  /// its parent has been popped from the region stack.
-  Optional<SourceMappingRegion> DeferredRegion;
-
   CounterExpressionBuilder Builder;
 
   /// A location in the most recently visited file or macro.
@@ -556,8 +543,11 @@ struct CounterCoverageMappingBuilder
   /// expressions cross file or macro boundaries.
   SourceLocation MostRecentLocation;
 
-  /// Location of the last terminated region.
-  Optional<std::pair<SourceLocation, size_t>> LastTerminatedRegion;
+  /// Whether the visitor at a terminate statement.
+  bool HasTerminateStmt = false;
+
+  /// Gap region counter after terminate statement.
+  Counter GapRegionCounter;
 
   /// Return a counter for the subtraction of \c RHS from \c LHS
   Counter subtractCounters(Counter LHS, Counter RHS) {
@@ -590,77 +580,13 @@ struct CounterCoverageMappingBuilder
 
     if (StartLoc && !FalseCount.hasValue()) {
       MostRecentLocation = *StartLoc;
-      completeDeferred(Count, MostRecentLocation);
     }
 
-    RegionStack.emplace_back(Count, FalseCount, StartLoc, EndLoc,
-                             FalseCount.hasValue());
+    RegionStack.emplace_back(Count, FalseCount, StartLoc, EndLoc);
 
     return RegionStack.size() - 1;
   }
 
-  /// Complete any pending deferred region by setting its end location and
-  /// count, and then pushing it onto the region stack.
-  size_t completeDeferred(Counter Count, SourceLocation DeferredEndLoc) {
-    size_t Index = RegionStack.size();
-    if (!DeferredRegion)
-      return Index;
-
-    // Consume the pending region.
-    SourceMappingRegion DR = DeferredRegion.getValue();
-    DeferredRegion = None;
-
-    // If the region ends in an expansion, find the expansion site.
-    FileID StartFile = SM.getFileID(DR.getBeginLoc());
-    if (SM.getFileID(DeferredEndLoc) != StartFile) {
-      if (isNestedIn(DeferredEndLoc, StartFile)) {
-        do {
-          DeferredEndLoc = getIncludeOrExpansionLoc(DeferredEndLoc);
-        } while (StartFile != SM.getFileID(DeferredEndLoc));
-      } else {
-        return Index;
-      }
-    }
-
-    // The parent of this deferred region ends where the containing decl ends,
-    // so the region isn't useful.
-    if (DR.getBeginLoc() == DeferredEndLoc)
-      return Index;
-
-    // If we're visiting statements in non-source order (e.g switch cases or
-    // a loop condition) we can't construct a sensible deferred region.
-    if (!SpellingRegion(SM, DR.getBeginLoc(), DeferredEndLoc).isInSourceOrder())
-      return Index;
-
-    DR.setGap(true);
-    DR.setCounter(Count);
-    DR.setEndLoc(DeferredEndLoc);
-    handleFileExit(DeferredEndLoc);
-    RegionStack.push_back(DR);
-    return Index;
-  }
-
-  /// Complete a deferred region created after a terminated region at the
-  /// top-level.
-  void completeTopLevelDeferredRegion(Counter Count,
-                                      SourceLocation DeferredEndLoc) {
-    if (DeferredRegion || !LastTerminatedRegion)
-      return;
-
-    if (LastTerminatedRegion->second != RegionStack.size())
-      return;
-
-    SourceLocation Start = LastTerminatedRegion->first;
-    if (SM.getFileID(Start) != SM.getMainFileID())
-      return;
-
-    SourceMappingRegion DR = RegionStack.back();
-    DR.setStartLoc(Start);
-    DR.setDeferred(false);
-    DeferredRegion = DR;
-    completeDeferred(Count, DeferredEndLoc);
-  }
-
   size_t locationDepth(SourceLocation Loc) {
     size_t Depth = 0;
     while (Loc.isValid()) {
@@ -676,7 +602,6 @@ struct CounterCoverageMappingBuilder
   /// function's \c SourceRegions.
   void popRegions(size_t ParentIndex) {
     assert(RegionStack.size() >= ParentIndex && "parent not in stack");
-    bool ParentOfDeferredRegion = false;
     while (RegionStack.size() > ParentIndex) {
       SourceMappingRegion &Region = RegionStack.back();
       if (Region.hasStartLoc()) {
@@ -746,32 +671,9 @@ struct CounterCoverageMappingBuilder
         assert(SM.isWrittenInSameFile(Region.getBeginLoc(), EndLoc));
         assert(SpellingRegion(SM, Region).isInSourceOrder());
         SourceRegions.push_back(Region);
-
-        if (ParentOfDeferredRegion) {
-          ParentOfDeferredRegion = false;
-
-          // If there's an existing deferred region, keep the old one, because
-          // it means there are two consecutive returns (or a similar pattern).
-          if (!DeferredRegion.hasValue() &&
-              // File IDs aren't gathered within macro expansions, so it isn't
-              // useful to try and create a deferred region inside of one.
-              !EndLoc.isMacroID())
-            DeferredRegion =
-                SourceMappingRegion(Counter::getZero(), EndLoc, None);
         }
-      } else if (Region.isDeferred()) {
-        assert(!ParentOfDeferredRegion && "Consecutive deferred regions");
-        ParentOfDeferredRegion = true;
-      }
       RegionStack.pop_back();
-
-      // If the zero region pushed after the last terminated region no longer
-      // exists, clear its cached information.
-      if (LastTerminatedRegion &&
-          RegionStack.size() < LastTerminatedRegion->second)
-        LastTerminatedRegion = None;
     }
-    assert(!ParentOfDeferredRegion && "Deferred region with no parent");
   }
 
   /// Return the currently active region.
@@ -955,8 +857,6 @@ struct CounterCoverageMappingBuilder
     handleFileExit(StartLoc);
     if (!Region.hasStartLoc())
       Region.setStartLoc(StartLoc);
-
-    completeDeferred(Region.getCounter(), StartLoc);
   }
 
   /// Mark \c S as a terminator, starting a zero region.
@@ -967,30 +867,56 @@ struct CounterCoverageMappingBuilder
     if (!Region.hasEndLoc())
       Region.setEndLoc(EndLoc);
     pushRegion(Counter::getZero());
-    auto &ZeroRegion = getRegion();
-    ZeroRegion.setDeferred(true);
-    LastTerminatedRegion = {EndLoc, RegionStack.size()};
+    HasTerminateStmt = true;
   }
 
   /// Find a valid gap range between \p AfterLoc and \p BeforeLoc.
   Optional<SourceRange> findGapAreaBetween(SourceLocation AfterLoc,
                                            SourceLocation BeforeLoc) {
+    // If AfterLoc is in function-like macro, use the right parenthesis
+    // location.
+    if (AfterLoc.isMacroID()) {
+      FileID FID = SM.getFileID(AfterLoc);
+      const SrcMgr::ExpansionInfo *EI = &SM.getSLocEntry(FID).getExpansion();
+      if (EI->isFunctionMacroExpansion())
+        AfterLoc = EI->getExpansionLocEnd();
+    }
+
+    size_t StartDepth = locationDepth(AfterLoc);
+    size_t EndDepth = locationDepth(BeforeLoc);
+    while (!SM.isWrittenInSameFile(AfterLoc, BeforeLoc)) {
+      bool UnnestStart = StartDepth >= EndDepth;
+      bool UnnestEnd = EndDepth >= StartDepth;
+      if (UnnestEnd) {
+        assert(SM.isWrittenInSameFile(getStartOfFileOrMacro(BeforeLoc),
+                                      BeforeLoc));
+
+        BeforeLoc = getIncludeOrExpansionLoc(BeforeLoc);
+        assert(BeforeLoc.isValid());
+        EndDepth--;
+      }
+      if (UnnestStart) {
+        assert(SM.isWrittenInSameFile(AfterLoc,
+                                      getEndOfFileOrMacro(AfterLoc)));
+
+        AfterLoc = getIncludeOrExpansionLoc(AfterLoc);
+        assert(AfterLoc.isValid());
+        AfterLoc = getPreciseTokenLocEnd(AfterLoc);
+        assert(AfterLoc.isValid());
+        StartDepth--;
+      }
+    }
+    AfterLoc = getPreciseTokenLocEnd(AfterLoc);
     // If the start and end locations of the gap are both within the same macro
     // file, the range may not be in source order.
     if (AfterLoc.isMacroID() || BeforeLoc.isMacroID())
       return None;
-    if (!SM.isWrittenInSameFile(AfterLoc, BeforeLoc))
+    if (!SM.isWrittenInSameFile(AfterLoc, BeforeLoc) ||
+        !SpellingRegion(SM, AfterLoc, BeforeLoc).isInSourceOrder())
       return None;
     return {{AfterLoc, BeforeLoc}};
   }
 
-  /// Find the source range after \p AfterStmt and before \p BeforeStmt.
-  Optional<SourceRange> findGapAreaBetween(const Stmt *AfterStmt,
-                                           const Stmt *BeforeStmt) {
-    return findGapAreaBetween(getPreciseTokenLocEnd(getEnd(AfterStmt)),
-                              getStart(BeforeStmt));
-  }
-
   /// Emit a gap region between \p StartLoc and \p EndLoc with the given count.
   void fillGapAreaWithCount(SourceLocation StartLoc, SourceLocation EndLoc,
                             Counter Count) {
@@ -1015,15 +941,13 @@ struct CounterCoverageMappingBuilder
       CoverageMappingModuleGen &CVM,
       llvm::DenseMap<const Stmt *, unsigned> &CounterMap, SourceManager &SM,
       const LangOptions &LangOpts)
-      : CoverageMappingBuilder(CVM, SM, LangOpts), CounterMap(CounterMap),
-        DeferredRegion(None) {}
+      : CoverageMappingBuilder(CVM, SM, LangOpts), CounterMap(CounterMap) {}
 
   /// Write the mapping data to the output stream
   void write(llvm::raw_ostream &OS) {
     llvm::SmallVector<unsigned, 8> VirtualFileMapping;
     gatherFileIDs(VirtualFileMapping);
     SourceRegionFilter Filter = emitExpansionRegions();
-    assert(!DeferredRegion && "Deferred region never completed");
     emitSourceRegions(Filter);
     gatherSkippedRegions();
 
@@ -1038,15 +962,32 @@ struct CounterCoverageMappingBuilder
   void VisitStmt(const Stmt *S) {
     if (S->getBeginLoc().isValid())
       extendRegion(S);
+    const Stmt *LastStmt = nullptr;
+    bool SaveTerminateStmt = HasTerminateStmt;
+    HasTerminateStmt = false;
+    GapRegionCounter = Counter::getZero();
     for (const Stmt *Child : S->children())
-      if (Child)
+      if (Child) {
+        // If last statement contains terminate statements, add a gap area
+        // between the two statements. Skipping attributed statements, because
+        // they don't have valid start location.
+        if (LastStmt && HasTerminateStmt && !dyn_cast<AttributedStmt>(Child)) {
+          auto Gap = findGapAreaBetween(getEnd(LastStmt), getStart(Child));
+          if (Gap)
+            fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(),
+                                 GapRegionCounter);
+          SaveTerminateStmt = true;
+          HasTerminateStmt = false;
+        }
         this->Visit(Child);
+        LastStmt = Child;
+      }
+    if (SaveTerminateStmt)
+      HasTerminateStmt = true;
     handleFileExit(getEnd(S));
   }
 
   void VisitDecl(const Decl *D) {
-    assert(!DeferredRegion && "Deferred region never completed");
-
     Stmt *Body = D->getBody();
 
     // Do not propagate region counts into system headers.
@@ -1064,11 +1005,6 @@ struct CounterCoverageMappingBuilder
     propagateCounts(getRegionCounter(Body), Body,
                     /*VisitChildren=*/!Defaulted);
     assert(RegionStack.empty() && "Regions entered but never exited");
-
-    // Discard the last uncompleted deferred region in a decl, if one exists.
-    // This prevents lines at the end of a function containing only whitespace
-    // or closing braces from being marked as uncovered.
-    DeferredRegion = None;
   }
 
   void VisitReturnStmt(const ReturnStmt *S) {
@@ -1102,8 +1038,6 @@ struct CounterCoverageMappingBuilder
   void VisitLabelStmt(const LabelStmt *S) {
     Counter LabelCount = getRegionCounter(S);
     SourceLocation Start = getStart(S);
-    completeTopLevelDeferredRegion(LabelCount, Start);
-    completeDeferred(LabelCount, Start);
     // We can't extendRegion here or we risk overlapping with our new region.
     handleFileExit(Start);
     pushRegion(LabelCount, Start);
@@ -1148,6 +1082,9 @@ struct CounterCoverageMappingBuilder
     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
     BreakContinue BC = BreakContinueStack.pop_back_val();
 
+    bool BodyHasTerminateStmt = HasTerminateStmt;
+    HasTerminateStmt = false;
+
     // Go back to handle the condition.
     Counter CondCount =
         addCounters(ParentCount, BackedgeCount, BC.ContinueCount);
@@ -1155,14 +1092,18 @@ struct CounterCoverageMappingBuilder
     adjustForOutOfOrderTraversal(getEnd(S));
 
     // The body count applies to the area immediately after the increment.
-    auto Gap = findGapAreaBetween(S->getCond(), S->getBody());
+    auto Gap = findGapAreaBetween(S->getRParenLoc(), getStart(S->getBody()));
     if (Gap)
       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
 
     Counter OutCount =
         addCounters(BC.BreakCount, subtractCounters(CondCount, BodyCount));
-    if (OutCount != ParentCount)
+    if (OutCount != ParentCount) {
       pushRegion(OutCount);
+      GapRegionCounter = OutCount;
+      if (BodyHasTerminateStmt)
+        HasTerminateStmt = true;
+    }
 
     // Create Branch Region around condition.
     createBranchRegion(S->getCond(), BodyCount,
@@ -1181,17 +1122,25 @@ struct CounterCoverageMappingBuilder
         propagateCounts(addCounters(ParentCount, BodyCount), S->getBody());
     BreakContinue BC = BreakContinueStack.pop_back_val();
 
+    bool BodyHasTerminateStmt = HasTerminateStmt;
+    HasTerminateStmt = false;
+
     Counter CondCount = addCounters(BackedgeCount, BC.ContinueCount);
     propagateCounts(CondCount, S->getCond());
 
     Counter OutCount =
         addCounters(BC.BreakCount, subtractCounters(CondCount, BodyCount));
-    if (OutCount != ParentCount)
+    if (OutCount != ParentCount) {
       pushRegion(OutCount);
+      GapRegionCounter = OutCount;
+    }
 
     // Create Branch Region around condition.
     createBranchRegion(S->getCond(), BodyCount,
                        subtractCounters(CondCount, BodyCount));
+
+    if (BodyHasTerminateStmt)
+      HasTerminateStmt = true;
   }
 
   void VisitForStmt(const ForStmt *S) {
@@ -1212,6 +1161,9 @@ struct CounterCoverageMappingBuilder
     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
     BreakContinue BodyBC = BreakContinueStack.pop_back_val();
 
+    bool BodyHasTerminateStmt = HasTerminateStmt;
+    HasTerminateStmt = false;
+
     // The increment is essentially part of the body but it needs to include
     // the count for all the continue statements.
     BreakContinue IncrementBC;
@@ -1230,15 +1182,18 @@ struct CounterCoverageMappingBuilder
     }
 
     // The body count applies to the area immediately after the increment.
-    auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
-                                  getStart(S->getBody()));
+    auto Gap = findGapAreaBetween(S->getRParenLoc(), getStart(S->getBody()));
     if (Gap)
       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
 
     Counter OutCount = addCounters(BodyBC.BreakCount, IncrementBC.BreakCount,
                                    subtractCounters(CondCount, BodyCount));
-    if (OutCount != ParentCount)
+    if (OutCount != ParentCount) {
       pushRegion(OutCount);
+      GapRegionCounter = OutCount;
+      if (BodyHasTerminateStmt)
+        HasTerminateStmt = true;
+    }
 
     // Create Branch Region around condition.
     createBranchRegion(S->getCond(), BodyCount,
@@ -1260,9 +1215,11 @@ struct CounterCoverageMappingBuilder
     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
     BreakContinue BC = BreakContinueStack.pop_back_val();
 
+    bool BodyHasTerminateStmt = HasTerminateStmt;
+    HasTerminateStmt = false;
+
     // The body count applies to the area immediately after the range.
-    auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
-                                  getStart(S->getBody()));
+    auto Gap = findGapAreaBetween(S->getRParenLoc(), getStart(S->getBody()));
     if (Gap)
       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
 
@@ -1270,8 +1227,12 @@ struct CounterCoverageMappingBuilder
         addCounters(ParentCount, BackedgeCount, BC.ContinueCount);
     Counter OutCount =
         addCounters(BC.BreakCount, subtractCounters(LoopCount, BodyCount));
-    if (OutCount != ParentCount)
+    if (OutCount != ParentCount) {
       pushRegion(OutCount);
+      GapRegionCounter = OutCount;
+      if (BodyHasTerminateStmt)
+        HasTerminateStmt = true;
+    }
 
     // Create Branch Region around condition.
     createBranchRegion(S->getCond(), BodyCount,
@@ -1291,8 +1252,7 @@ struct CounterCoverageMappingBuilder
     BreakContinue BC = BreakContinueStack.pop_back_val();
 
     // The body count applies to the area immediately after the collection.
-    auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
-                                  getStart(S->getBody()));
+    auto Gap = findGapAreaBetween(S->getRParenLoc(), getStart(S->getBody()));
     if (Gap)
       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
 
@@ -1300,8 +1260,10 @@ struct CounterCoverageMappingBuilder
         addCounters(ParentCount, BackedgeCount, BC.ContinueCount);
     Counter OutCount =
         addCounters(BC.BreakCount, subtractCounters(LoopCount, BodyCount));
-    if (OutCount != ParentCount)
+    if (OutCount != ParentCount) {
       pushRegion(OutCount);
+      GapRegionCounter = OutCount;
+    }
   }
 
   void VisitSwitchStmt(const SwitchStmt *S) {
@@ -1321,8 +1283,7 @@ struct CounterCoverageMappingBuilder
         // the unreachable code at the beginning of the switch body.
         size_t Index = pushRegion(Counter::getZero(), getStart(CS));
         getRegion().setGap(true);
-        for (const auto *Child : CS->children())
-          Visit(Child);
+        Visit(Body);
 
         // Set the end for the body of the switch, if it isn't already set.
         for (size_t i = RegionStack.size(); i != Index; --i) {
@@ -1344,6 +1305,7 @@ struct CounterCoverageMappingBuilder
     Counter ExitCount = getRegionCounter(S);
     SourceLocation ExitLoc = getEnd(S);
     pushRegion(ExitCount);
+    GapRegionCounter = ExitCount;
 
     // Ensure that handleFileExit recognizes when the end location is located
     // in a different file.
@@ -1386,6 +1348,8 @@ struct CounterCoverageMappingBuilder
     else
       pushRegion(Count, getStart(S));
 
+    GapRegionCounter = Count;
+
     if (const auto *CS = dyn_cast<CaseStmt>(S)) {
       Visit(CS->getLHS());
       if (const Expr *RHS = CS->getRHS())
@@ -1411,7 +1375,7 @@ struct CounterCoverageMappingBuilder
     propagateCounts(ParentCount, S->getCond());
 
     // The 'then' count applies to the area immediately after the condition.
-    auto Gap = findGapAreaBetween(S->getCond(), S->getThen());
+    auto Gap = findGapAreaBetween(S->getRParenLoc(), getStart(S->getThen()));
     if (Gap)
       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), ThenCount);
 
@@ -1420,17 +1384,25 @@ struct CounterCoverageMappingBuilder
 
     Counter ElseCount = subtractCounters(ParentCount, ThenCount);
     if (const Stmt *Else = S->getElse()) {
+      bool ThenHasTerminateStmt = HasTerminateStmt;
+      HasTerminateStmt = false;
+
       // The 'else' count applies to the area immediately after the 'then'.
-      Gap = findGapAreaBetween(S->getThen(), Else);
+      Gap = findGapAreaBetween(getEnd(S->getThen()), getStart(Else));
       if (Gap)
         fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), ElseCount);
       extendRegion(Else);
       OutCount = addCounters(OutCount, propagateCounts(ElseCount, Else));
+
+      if (ThenHasTerminateStmt)
+        HasTerminateStmt = true;
     } else
       OutCount = addCounters(OutCount, ElseCount);
 
-    if (OutCount != ParentCount)
+    if (OutCount != ParentCount) {
       pushRegion(OutCount);
+      GapRegionCounter = OutCount;
+    }
 
     // Create Branch Region around condition.
     createBranchRegion(S->getCond(), ThenCount,
@@ -1462,7 +1434,7 @@ struct CounterCoverageMappingBuilder
     Counter ParentCount = getRegion().getCounter();
     Counter TrueCount = getRegionCounter(E);
 
-    Visit(E->getCond());
+    propagateCounts(ParentCount, E->getCond());
 
     if (!isa<BinaryConditionalOperator>(E)) {
       // The 'then' count applies to the area immediately after the condition.
@@ -1588,14 +1560,22 @@ static void dump(llvm::raw_ostream &OS, StringRef FunctionName,
 CoverageMappingModuleGen::CoverageMappingModuleGen(
     CodeGenModule &CGM, CoverageSourceInfo &SourceInfo)
     : CGM(CGM), SourceInfo(SourceInfo) {
-  ProfilePrefixMap = CGM.getCodeGenOpts().ProfilePrefixMap;
+  CoveragePrefixMap = CGM.getCodeGenOpts().CoveragePrefixMap;
+}
+
+std::string CoverageMappingModuleGen::getCurrentDirname() {
+  if (!CGM.getCodeGenOpts().CoverageCompilationDir.empty())
+    return CGM.getCodeGenOpts().CoverageCompilationDir;
+
+  SmallString<256> CWD;
+  llvm::sys::fs::current_path(CWD);
+  return CWD.str().str();
 }
 
 std::string CoverageMappingModuleGen::normalizeFilename(StringRef Filename) {
   llvm::SmallString<256> Path(Filename);
-  llvm::sys::fs::make_absolute(Path);
   llvm::sys::path::remove_dots(Path, /*remove_dot_dot=*/true);
-  for (const auto &Entry : ProfilePrefixMap) {
+  for (const auto &Entry : CoveragePrefixMap) {
     if (llvm::sys::path::replace_path_prefix(Path, Entry.first, Entry.second))
       break;
   }
@@ -1674,18 +1654,17 @@ void CoverageMappingModuleGen::addFunctionMappingRecord(
     // also processed by the CoverageMappingWriter which performs
     // additional minimization operations such as reducing the number of
     // expressions.
+    llvm::SmallVector<std::string, 16> FilenameStrs;
     std::vector<StringRef> Filenames;
     std::vector<CounterExpression> Expressions;
     std::vector<CounterMappingRegion> Regions;
-    llvm::SmallVector<std::string, 16> FilenameStrs;
-    llvm::SmallVector<StringRef, 16> FilenameRefs;
-    FilenameStrs.resize(FileEntries.size());
-    FilenameRefs.resize(FileEntries.size());
+    FilenameStrs.resize(FileEntries.size() + 1);
+    FilenameStrs[0] = normalizeFilename(getCurrentDirname());
     for (const auto &Entry : FileEntries) {
       auto I = Entry.second;
       FilenameStrs[I] = normalizeFilename(Entry.first->getName());
-      FilenameRefs[I] = FilenameStrs[I];
     }
+    ArrayRef<std::string> FilenameRefs = llvm::makeArrayRef(FilenameStrs);
     RawCoverageMappingReader Reader(CoverageMapping, FilenameRefs, Filenames,
                                     Expressions, Regions);
     if (Reader.read())
@@ -1702,19 +1681,18 @@ void CoverageMappingModuleGen::emit() {
 
   // Create the filenames and merge them with coverage mappings
   llvm::SmallVector<std::string, 16> FilenameStrs;
-  llvm::SmallVector<StringRef, 16> FilenameRefs;
-  FilenameStrs.resize(FileEntries.size());
-  FilenameRefs.resize(FileEntries.size());
+  FilenameStrs.resize(FileEntries.size() + 1);
+  // The first filename is the current working directory.
+  FilenameStrs[0] = normalizeFilename(getCurrentDirname());
   for (const auto &Entry : FileEntries) {
     auto I = Entry.second;
     FilenameStrs[I] = normalizeFilename(Entry.first->getName());
-    FilenameRefs[I] = FilenameStrs[I];
   }
 
   std::string Filenames;
   {
     llvm::raw_string_ostream OS(Filenames);
-    CoverageFilenamesSectionWriter(FilenameRefs).write(OS);
+    CoverageFilenamesSectionWriter(FilenameStrs).write(OS);
   }
   auto *FilenamesVal =
       llvm::ConstantDataArray::getString(Ctx, Filenames, false);
@@ -1772,7 +1750,7 @@ unsigned CoverageMappingModuleGen::getFileID(const FileEntry *File) {
   auto It = FileEntries.find(File);
   if (It != FileEntries.end())
     return It->second;
-  unsigned FileID = FileEntries.size();
+  unsigned FileID = FileEntries.size() + 1;
   FileEntries.insert(std::make_pair(File, FileID));
   return FileID;
 }