1 files changed, 86 insertions, 5 deletions

diff --git a/clang/lib/CodeGen/CoverageMappingGen.cpp b/clang/lib/CodeGen/CoverageMappingGen.cpp
index 28aabd0f071..27d98dc60d7 100644
--- a/clang/lib/CodeGen/CoverageMappingGen.cpp
+++ b/clang/lib/CodeGen/CoverageMappingGen.cpp
@@ -45,10 +45,14 @@ class SourceMappingRegion {
   /// \brief The region's ending location.
   Optional<SourceLocation> LocEnd;
 
+  /// Whether this region should be emitted after its parent is emitted.
+  bool DeferRegion;
+
 public:
   SourceMappingRegion(Counter Count, Optional<SourceLocation> LocStart,
-                      Optional<SourceLocation> LocEnd)
-      : Count(Count), LocStart(LocStart), LocEnd(LocEnd) {}
+                      Optional<SourceLocation> LocEnd, bool DeferRegion = false)
+      : Count(Count), LocStart(LocStart), LocEnd(LocEnd),
+        DeferRegion(DeferRegion) {}
 
   const Counter &getCounter() const { return Count; }
 
@@ -71,6 +75,10 @@ public:
     assert(LocEnd && "Region has no end location");
     return *LocEnd;
   }
+
+  bool isDeferred() const { return DeferRegion; }
+
+  void setDeferred(bool Deferred) { DeferRegion = Deferred; }
 };
 
 /// Spelling locations for the start and end of a source region.
@@ -409,6 +417,10 @@ struct CounterCoverageMappingBuilder
   /// \brief 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;
 
   /// \brief A location in the most recently visited file or macro.
@@ -444,19 +456,60 @@ struct CounterCoverageMappingBuilder
   /// used with popRegions to exit a "scope", ending the region that was pushed.
   size_t pushRegion(Counter Count, Optional<SourceLocation> StartLoc = None,
                     Optional<SourceLocation> EndLoc = None) {
-    if (StartLoc)
+    if (StartLoc) {
       MostRecentLocation = *StartLoc;
+      completeDeferred(Count, MostRecentLocation);
+    }
     RegionStack.emplace_back(Count, 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.
+    if (SM.getFileID(DeferredEndLoc) != SM.getMainFileID()) {
+      FileID StartFile = SM.getFileID(DR.getStartLoc());
+      if (isNestedIn(DeferredEndLoc, StartFile)) {
+        do {
+          DeferredEndLoc = getIncludeOrExpansionLoc(DeferredEndLoc);
+        } while (StartFile != SM.getFileID(DeferredEndLoc));
+      }
+    }
+
+    // The parent of this deferred region ends where the containing decl ends,
+    // so the region isn't useful.
+    if (DR.getStartLoc() == 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.getStartLoc(), DeferredEndLoc).isInSourceOrder())
+      return Index;
+
+    DR.setCounter(Count);
+    DR.setEndLoc(DeferredEndLoc);
+    handleFileExit(DeferredEndLoc);
+    RegionStack.push_back(DR);
+    return Index;
+  }
+
   /// \brief Pop regions from the stack into the function's list of regions.
   ///
   /// Adds all regions from \c ParentIndex to the top of the stack to the
   /// 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()) {
@@ -488,9 +541,26 @@ struct CounterCoverageMappingBuilder
 
         assert(SM.isWrittenInSameFile(Region.getStartLoc(), EndLoc));
         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.
+              (SM.getFileID(EndLoc) == SM.getMainFileID()))
+            DeferredRegion =
+                SourceMappingRegion(Counter::getZero(), EndLoc, None);
+        }
+      } else if (Region.isDeferred()) {
+        assert(!ParentOfDeferredRegion && "Consecutive deferred regions");
+        ParentOfDeferredRegion = true;
       }
       RegionStack.pop_back();
     }
+    assert(!ParentOfDeferredRegion && "Deferred region with no parent");
   }
 
   /// \brief Return the currently active region.
@@ -617,6 +687,8 @@ struct CounterCoverageMappingBuilder
     handleFileExit(StartLoc);
     if (!Region.hasStartLoc())
       Region.setStartLoc(StartLoc);
+
+    completeDeferred(Region.getCounter(), StartLoc);
   }
 
   /// \brief Mark \c S as a terminator, starting a zero region.
@@ -626,6 +698,7 @@ struct CounterCoverageMappingBuilder
     if (!Region.hasEndLoc())
       Region.setEndLoc(getEnd(S));
     pushRegion(Counter::getZero());
+    getRegion().setDeferred(true);
   }
 
   /// \brief Keep counts of breaks and continues inside loops.
@@ -639,13 +712,15 @@ struct CounterCoverageMappingBuilder
       CoverageMappingModuleGen &CVM,
       llvm::DenseMap<const Stmt *, unsigned> &CounterMap, SourceManager &SM,
       const LangOptions &LangOpts)
-      : CoverageMappingBuilder(CVM, SM, LangOpts), CounterMap(CounterMap) {}
+      : CoverageMappingBuilder(CVM, SM, LangOpts), CounterMap(CounterMap),
+        DeferredRegion(None) {}
 
   /// \brief 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();
 
@@ -667,13 +742,19 @@ struct CounterCoverageMappingBuilder
   }
 
   void VisitDecl(const Decl *D) {
+    assert(!DeferredRegion && "Deferred region never completed");
+
     Stmt *Body = D->getBody();
 
     // Do not propagate region counts into system headers.
     if (Body && SM.isInSystemHeader(SM.getSpellingLoc(getStart(Body))))
       return;
 
-    propagateCounts(getRegionCounter(Body), Body);
+    Counter ExitCount = propagateCounts(getRegionCounter(Body), Body);
+    assert(RegionStack.empty() && "Regions entered but never exited");
+
+    // Complete any deferred regions introduced by the last statement in a decl.
+    popRegions(completeDeferred(ExitCount, getEnd(Body)));
   }
 
   void VisitReturnStmt(const ReturnStmt *S) {