1 //===--- CoverageMappingGen.cpp - Coverage mapping generation ---*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Instrumentation-based code coverage mapping generator
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "CoverageMappingGen.h"
14 #include "CodeGenFunction.h"
15 #include "clang/AST/StmtVisitor.h"
16 #include "clang/Basic/Diagnostic.h"
17 #include "clang/Basic/FileManager.h"
18 #include "clang/Frontend/FrontendDiagnostic.h"
19 #include "clang/Lex/Lexer.h"
20 #include "llvm/ADT/Optional.h"
21 #include "llvm/ADT/SmallSet.h"
22 #include "llvm/ADT/StringExtras.h"
23 #include "llvm/ProfileData/Coverage/CoverageMapping.h"
24 #include "llvm/ProfileData/Coverage/CoverageMappingReader.h"
25 #include "llvm/ProfileData/Coverage/CoverageMappingWriter.h"
26 #include "llvm/ProfileData/InstrProfReader.h"
27 #include "llvm/Support/FileSystem.h"
28 #include "llvm/Support/Path.h"
29 
30 // This selects the coverage mapping format defined when `InstrProfData.inc`
31 // is textually included.
32 #define COVMAP_V3
33 
34 using namespace clang;
35 using namespace CodeGen;
36 using namespace llvm::coverage;
37 
38 CoverageSourceInfo *
39 CoverageMappingModuleGen::setUpCoverageCallbacks(Preprocessor &PP) {
40   CoverageSourceInfo *CoverageInfo = new CoverageSourceInfo();
41   PP.addPPCallbacks(std::unique_ptr<PPCallbacks>(CoverageInfo));
42   PP.addCommentHandler(CoverageInfo);
43   PP.setPreprocessToken(true);
44   PP.setTokenWatcher([CoverageInfo](clang::Token Tok) {
45     // Update previous token location.
46     CoverageInfo->PrevTokLoc = Tok.getLocation();
47     CoverageInfo->updateNextTokLoc(Tok.getLocation());
48   });
49   return CoverageInfo;
50 }
51 
52 void CoverageSourceInfo::SourceRangeSkipped(SourceRange Range, SourceLocation) {
53   SkippedRanges.push_back({Range});
54 }
55 
56 bool CoverageSourceInfo::HandleComment(Preprocessor &PP, SourceRange Range) {
57   SkippedRanges.push_back({Range, PrevTokLoc});
58   AfterComment = true;
59   return false;
60 }
61 
62 void CoverageSourceInfo::updateNextTokLoc(SourceLocation Loc) {
63   if (AfterComment) {
64     SkippedRanges.back().NextTokLoc = Loc;
65     AfterComment = false;
66   }
67 }
68 
69 namespace {
70 
71 /// A region of source code that can be mapped to a counter.
72 class SourceMappingRegion {
73   Counter Count;
74 
75   /// The region's starting location.
76   Optional<SourceLocation> LocStart;
77 
78   /// The region's ending location.
79   Optional<SourceLocation> LocEnd;
80 
81   /// Whether this region should be emitted after its parent is emitted.
82   bool DeferRegion;
83 
84   /// Whether this region is a gap region. The count from a gap region is set
85   /// as the line execution count if there are no other regions on the line.
86   bool GapRegion;
87 
88 public:
89   SourceMappingRegion(Counter Count, Optional<SourceLocation> LocStart,
90                       Optional<SourceLocation> LocEnd, bool DeferRegion = false,
91                       bool GapRegion = false)
92       : Count(Count), LocStart(LocStart), LocEnd(LocEnd),
93         DeferRegion(DeferRegion), GapRegion(GapRegion) {}
94 
95   const Counter &getCounter() const { return Count; }
96 
97   void setCounter(Counter C) { Count = C; }
98 
99   bool hasStartLoc() const { return LocStart.hasValue(); }
100 
101   void setStartLoc(SourceLocation Loc) { LocStart = Loc; }
102 
103   SourceLocation getBeginLoc() const {
104     assert(LocStart && "Region has no start location");
105     return *LocStart;
106   }
107 
108   bool hasEndLoc() const { return LocEnd.hasValue(); }
109 
110   void setEndLoc(SourceLocation Loc) {
111     assert(Loc.isValid() && "Setting an invalid end location");
112     LocEnd = Loc;
113   }
114 
115   SourceLocation getEndLoc() const {
116     assert(LocEnd && "Region has no end location");
117     return *LocEnd;
118   }
119 
120   bool isDeferred() const { return DeferRegion; }
121 
122   void setDeferred(bool Deferred) { DeferRegion = Deferred; }
123 
124   bool isGap() const { return GapRegion; }
125 
126   void setGap(bool Gap) { GapRegion = Gap; }
127 };
128 
129 /// Spelling locations for the start and end of a source region.
130 struct SpellingRegion {
131   /// The line where the region starts.
132   unsigned LineStart;
133 
134   /// The column where the region starts.
135   unsigned ColumnStart;
136 
137   /// The line where the region ends.
138   unsigned LineEnd;
139 
140   /// The column where the region ends.
141   unsigned ColumnEnd;
142 
143   SpellingRegion(SourceManager &SM, SourceLocation LocStart,
144                  SourceLocation LocEnd) {
145     LineStart = SM.getSpellingLineNumber(LocStart);
146     ColumnStart = SM.getSpellingColumnNumber(LocStart);
147     LineEnd = SM.getSpellingLineNumber(LocEnd);
148     ColumnEnd = SM.getSpellingColumnNumber(LocEnd);
149   }
150 
151   SpellingRegion(SourceManager &SM, SourceMappingRegion &R)
152       : SpellingRegion(SM, R.getBeginLoc(), R.getEndLoc()) {}
153 
154   /// Check if the start and end locations appear in source order, i.e
155   /// top->bottom, left->right.
156   bool isInSourceOrder() const {
157     return (LineStart < LineEnd) ||
158            (LineStart == LineEnd && ColumnStart <= ColumnEnd);
159   }
160 };
161 
162 /// Provides the common functionality for the different
163 /// coverage mapping region builders.
164 class CoverageMappingBuilder {
165 public:
166   CoverageMappingModuleGen &CVM;
167   SourceManager &SM;
168   const LangOptions &LangOpts;
169 
170 private:
171   /// Map of clang's FileIDs to IDs used for coverage mapping.
172   llvm::SmallDenseMap<FileID, std::pair<unsigned, SourceLocation>, 8>
173       FileIDMapping;
174 
175 public:
176   /// The coverage mapping regions for this function
177   llvm::SmallVector<CounterMappingRegion, 32> MappingRegions;
178   /// The source mapping regions for this function.
179   std::vector<SourceMappingRegion> SourceRegions;
180 
181   /// A set of regions which can be used as a filter.
182   ///
183   /// It is produced by emitExpansionRegions() and is used in
184   /// emitSourceRegions() to suppress producing code regions if
185   /// the same area is covered by expansion regions.
186   typedef llvm::SmallSet<std::pair<SourceLocation, SourceLocation>, 8>
187       SourceRegionFilter;
188 
189   CoverageMappingBuilder(CoverageMappingModuleGen &CVM, SourceManager &SM,
190                          const LangOptions &LangOpts)
191       : CVM(CVM), SM(SM), LangOpts(LangOpts) {}
192 
193   /// Return the precise end location for the given token.
194   SourceLocation getPreciseTokenLocEnd(SourceLocation Loc) {
195     // We avoid getLocForEndOfToken here, because it doesn't do what we want for
196     // macro locations, which we just treat as expanded files.
197     unsigned TokLen =
198         Lexer::MeasureTokenLength(SM.getSpellingLoc(Loc), SM, LangOpts);
199     return Loc.getLocWithOffset(TokLen);
200   }
201 
202   /// Return the start location of an included file or expanded macro.
203   SourceLocation getStartOfFileOrMacro(SourceLocation Loc) {
204     if (Loc.isMacroID())
205       return Loc.getLocWithOffset(-SM.getFileOffset(Loc));
206     return SM.getLocForStartOfFile(SM.getFileID(Loc));
207   }
208 
209   /// Return the end location of an included file or expanded macro.
210   SourceLocation getEndOfFileOrMacro(SourceLocation Loc) {
211     if (Loc.isMacroID())
212       return Loc.getLocWithOffset(SM.getFileIDSize(SM.getFileID(Loc)) -
213                                   SM.getFileOffset(Loc));
214     return SM.getLocForEndOfFile(SM.getFileID(Loc));
215   }
216 
217   /// Find out where the current file is included or macro is expanded.
218   SourceLocation getIncludeOrExpansionLoc(SourceLocation Loc) {
219     return Loc.isMacroID() ? SM.getImmediateExpansionRange(Loc).getBegin()
220                            : SM.getIncludeLoc(SM.getFileID(Loc));
221   }
222 
223   /// Return true if \c Loc is a location in a built-in macro.
224   bool isInBuiltin(SourceLocation Loc) {
225     return SM.getBufferName(SM.getSpellingLoc(Loc)) == "<built-in>";
226   }
227 
228   /// Check whether \c Loc is included or expanded from \c Parent.
229   bool isNestedIn(SourceLocation Loc, FileID Parent) {
230     do {
231       Loc = getIncludeOrExpansionLoc(Loc);
232       if (Loc.isInvalid())
233         return false;
234     } while (!SM.isInFileID(Loc, Parent));
235     return true;
236   }
237 
238   /// Get the start of \c S ignoring macro arguments and builtin macros.
239   SourceLocation getStart(const Stmt *S) {
240     SourceLocation Loc = S->getBeginLoc();
241     while (SM.isMacroArgExpansion(Loc) || isInBuiltin(Loc))
242       Loc = SM.getImmediateExpansionRange(Loc).getBegin();
243     return Loc;
244   }
245 
246   /// Get the end of \c S ignoring macro arguments and builtin macros.
247   SourceLocation getEnd(const Stmt *S) {
248     SourceLocation Loc = S->getEndLoc();
249     while (SM.isMacroArgExpansion(Loc) || isInBuiltin(Loc))
250       Loc = SM.getImmediateExpansionRange(Loc).getBegin();
251     return getPreciseTokenLocEnd(Loc);
252   }
253 
254   /// Find the set of files we have regions for and assign IDs
255   ///
256   /// Fills \c Mapping with the virtual file mapping needed to write out
257   /// coverage and collects the necessary file information to emit source and
258   /// expansion regions.
259   void gatherFileIDs(SmallVectorImpl<unsigned> &Mapping) {
260     FileIDMapping.clear();
261 
262     llvm::SmallSet<FileID, 8> Visited;
263     SmallVector<std::pair<SourceLocation, unsigned>, 8> FileLocs;
264     for (const auto &Region : SourceRegions) {
265       SourceLocation Loc = Region.getBeginLoc();
266       FileID File = SM.getFileID(Loc);
267       if (!Visited.insert(File).second)
268         continue;
269 
270       // Do not map FileID's associated with system headers.
271       if (SM.isInSystemHeader(SM.getSpellingLoc(Loc)))
272         continue;
273 
274       unsigned Depth = 0;
275       for (SourceLocation Parent = getIncludeOrExpansionLoc(Loc);
276            Parent.isValid(); Parent = getIncludeOrExpansionLoc(Parent))
277         ++Depth;
278       FileLocs.push_back(std::make_pair(Loc, Depth));
279     }
280     llvm::stable_sort(FileLocs, llvm::less_second());
281 
282     for (const auto &FL : FileLocs) {
283       SourceLocation Loc = FL.first;
284       FileID SpellingFile = SM.getDecomposedSpellingLoc(Loc).first;
285       auto Entry = SM.getFileEntryForID(SpellingFile);
286       if (!Entry)
287         continue;
288 
289       FileIDMapping[SM.getFileID(Loc)] = std::make_pair(Mapping.size(), Loc);
290       Mapping.push_back(CVM.getFileID(Entry));
291     }
292   }
293 
294   /// Get the coverage mapping file ID for \c Loc.
295   ///
296   /// If such file id doesn't exist, return None.
297   Optional<unsigned> getCoverageFileID(SourceLocation Loc) {
298     auto Mapping = FileIDMapping.find(SM.getFileID(Loc));
299     if (Mapping != FileIDMapping.end())
300       return Mapping->second.first;
301     return None;
302   }
303 
304   /// This shrinks the skipped range if it spans a line that contains a
305   /// non-comment token. If shrinking the skipped range would make it empty,
306   /// this returns None.
307   Optional<SpellingRegion> adjustSkippedRange(SourceManager &SM,
308                                               SpellingRegion SR,
309                                               SourceLocation PrevTokLoc,
310                                               SourceLocation NextTokLoc) {
311     // If Range begin location is invalid, it's not a comment region.
312     if (PrevTokLoc.isInvalid())
313       return SR;
314     unsigned PrevTokLine = SM.getSpellingLineNumber(PrevTokLoc);
315     unsigned NextTokLine = SM.getSpellingLineNumber(NextTokLoc);
316     SpellingRegion newSR(SR);
317     if (SR.LineStart == PrevTokLine) {
318       newSR.LineStart = SR.LineStart + 1;
319       newSR.ColumnStart = 1;
320     }
321     if (SR.LineEnd == NextTokLine) {
322       newSR.LineEnd = SR.LineEnd - 1;
323       newSR.ColumnEnd = SR.ColumnStart + 1;
324     }
325     if (newSR.isInSourceOrder())
326       return newSR;
327     return None;
328   }
329 
330   /// Gather all the regions that were skipped by the preprocessor
331   /// using the constructs like #if or comments.
332   void gatherSkippedRegions() {
333     /// An array of the minimum lineStarts and the maximum lineEnds
334     /// for mapping regions from the appropriate source files.
335     llvm::SmallVector<std::pair<unsigned, unsigned>, 8> FileLineRanges;
336     FileLineRanges.resize(
337         FileIDMapping.size(),
338         std::make_pair(std::numeric_limits<unsigned>::max(), 0));
339     for (const auto &R : MappingRegions) {
340       FileLineRanges[R.FileID].first =
341           std::min(FileLineRanges[R.FileID].first, R.LineStart);
342       FileLineRanges[R.FileID].second =
343           std::max(FileLineRanges[R.FileID].second, R.LineEnd);
344     }
345 
346     auto SkippedRanges = CVM.getSourceInfo().getSkippedRanges();
347     for (auto &I : SkippedRanges) {
348       SourceRange Range = I.Range;
349       auto LocStart = Range.getBegin();
350       auto LocEnd = Range.getEnd();
351       assert(SM.isWrittenInSameFile(LocStart, LocEnd) &&
352              "region spans multiple files");
353 
354       auto CovFileID = getCoverageFileID(LocStart);
355       if (!CovFileID)
356         continue;
357       SpellingRegion SR{SM, LocStart, LocEnd};
358       if (Optional<SpellingRegion> res =
359               adjustSkippedRange(SM, SR, I.PrevTokLoc, I.NextTokLoc))
360         SR = res.getValue();
361       else
362         continue;
363       auto Region = CounterMappingRegion::makeSkipped(
364           *CovFileID, SR.LineStart, SR.ColumnStart, SR.LineEnd, SR.ColumnEnd);
365       // Make sure that we only collect the regions that are inside
366       // the source code of this function.
367       if (Region.LineStart >= FileLineRanges[*CovFileID].first &&
368           Region.LineEnd <= FileLineRanges[*CovFileID].second)
369         MappingRegions.push_back(Region);
370     }
371   }
372 
373   /// Generate the coverage counter mapping regions from collected
374   /// source regions.
375   void emitSourceRegions(const SourceRegionFilter &Filter) {
376     for (const auto &Region : SourceRegions) {
377       assert(Region.hasEndLoc() && "incomplete region");
378 
379       SourceLocation LocStart = Region.getBeginLoc();
380       assert(SM.getFileID(LocStart).isValid() && "region in invalid file");
381 
382       // Ignore regions from system headers.
383       if (SM.isInSystemHeader(SM.getSpellingLoc(LocStart)))
384         continue;
385 
386       auto CovFileID = getCoverageFileID(LocStart);
387       // Ignore regions that don't have a file, such as builtin macros.
388       if (!CovFileID)
389         continue;
390 
391       SourceLocation LocEnd = Region.getEndLoc();
392       assert(SM.isWrittenInSameFile(LocStart, LocEnd) &&
393              "region spans multiple files");
394 
395       // Don't add code regions for the area covered by expansion regions.
396       // This not only suppresses redundant regions, but sometimes prevents
397       // creating regions with wrong counters if, for example, a statement's
398       // body ends at the end of a nested macro.
399       if (Filter.count(std::make_pair(LocStart, LocEnd)))
400         continue;
401 
402       // Find the spelling locations for the mapping region.
403       SpellingRegion SR{SM, LocStart, LocEnd};
404       assert(SR.isInSourceOrder() && "region start and end out of order");
405 
406       if (Region.isGap()) {
407         MappingRegions.push_back(CounterMappingRegion::makeGapRegion(
408             Region.getCounter(), *CovFileID, SR.LineStart, SR.ColumnStart,
409             SR.LineEnd, SR.ColumnEnd));
410       } else {
411         MappingRegions.push_back(CounterMappingRegion::makeRegion(
412             Region.getCounter(), *CovFileID, SR.LineStart, SR.ColumnStart,
413             SR.LineEnd, SR.ColumnEnd));
414       }
415     }
416   }
417 
418   /// Generate expansion regions for each virtual file we've seen.
419   SourceRegionFilter emitExpansionRegions() {
420     SourceRegionFilter Filter;
421     for (const auto &FM : FileIDMapping) {
422       SourceLocation ExpandedLoc = FM.second.second;
423       SourceLocation ParentLoc = getIncludeOrExpansionLoc(ExpandedLoc);
424       if (ParentLoc.isInvalid())
425         continue;
426 
427       auto ParentFileID = getCoverageFileID(ParentLoc);
428       if (!ParentFileID)
429         continue;
430       auto ExpandedFileID = getCoverageFileID(ExpandedLoc);
431       assert(ExpandedFileID && "expansion in uncovered file");
432 
433       SourceLocation LocEnd = getPreciseTokenLocEnd(ParentLoc);
434       assert(SM.isWrittenInSameFile(ParentLoc, LocEnd) &&
435              "region spans multiple files");
436       Filter.insert(std::make_pair(ParentLoc, LocEnd));
437 
438       SpellingRegion SR{SM, ParentLoc, LocEnd};
439       assert(SR.isInSourceOrder() && "region start and end out of order");
440       MappingRegions.push_back(CounterMappingRegion::makeExpansion(
441           *ParentFileID, *ExpandedFileID, SR.LineStart, SR.ColumnStart,
442           SR.LineEnd, SR.ColumnEnd));
443     }
444     return Filter;
445   }
446 };
447 
448 /// Creates unreachable coverage regions for the functions that
449 /// are not emitted.
450 struct EmptyCoverageMappingBuilder : public CoverageMappingBuilder {
451   EmptyCoverageMappingBuilder(CoverageMappingModuleGen &CVM, SourceManager &SM,
452                               const LangOptions &LangOpts)
453       : CoverageMappingBuilder(CVM, SM, LangOpts) {}
454 
455   void VisitDecl(const Decl *D) {
456     if (!D->hasBody())
457       return;
458     auto Body = D->getBody();
459     SourceLocation Start = getStart(Body);
460     SourceLocation End = getEnd(Body);
461     if (!SM.isWrittenInSameFile(Start, End)) {
462       // Walk up to find the common ancestor.
463       // Correct the locations accordingly.
464       FileID StartFileID = SM.getFileID(Start);
465       FileID EndFileID = SM.getFileID(End);
466       while (StartFileID != EndFileID && !isNestedIn(End, StartFileID)) {
467         Start = getIncludeOrExpansionLoc(Start);
468         assert(Start.isValid() &&
469                "Declaration start location not nested within a known region");
470         StartFileID = SM.getFileID(Start);
471       }
472       while (StartFileID != EndFileID) {
473         End = getPreciseTokenLocEnd(getIncludeOrExpansionLoc(End));
474         assert(End.isValid() &&
475                "Declaration end location not nested within a known region");
476         EndFileID = SM.getFileID(End);
477       }
478     }
479     SourceRegions.emplace_back(Counter(), Start, End);
480   }
481 
482   /// Write the mapping data to the output stream
483   void write(llvm::raw_ostream &OS) {
484     SmallVector<unsigned, 16> FileIDMapping;
485     gatherFileIDs(FileIDMapping);
486     emitSourceRegions(SourceRegionFilter());
487 
488     if (MappingRegions.empty())
489       return;
490 
491     CoverageMappingWriter Writer(FileIDMapping, None, MappingRegions);
492     Writer.write(OS);
493   }
494 };
495 
496 /// A StmtVisitor that creates coverage mapping regions which map
497 /// from the source code locations to the PGO counters.
498 struct CounterCoverageMappingBuilder
499     : public CoverageMappingBuilder,
500       public ConstStmtVisitor<CounterCoverageMappingBuilder> {
501   /// The map of statements to count values.
502   llvm::DenseMap<const Stmt *, unsigned> &CounterMap;
503 
504   /// A stack of currently live regions.
505   std::vector<SourceMappingRegion> RegionStack;
506 
507   /// The currently deferred region: its end location and count can be set once
508   /// its parent has been popped from the region stack.
509   Optional<SourceMappingRegion> DeferredRegion;
510 
511   CounterExpressionBuilder Builder;
512 
513   /// A location in the most recently visited file or macro.
514   ///
515   /// This is used to adjust the active source regions appropriately when
516   /// expressions cross file or macro boundaries.
517   SourceLocation MostRecentLocation;
518 
519   /// Location of the last terminated region.
520   Optional<std::pair<SourceLocation, size_t>> LastTerminatedRegion;
521 
522   /// Return a counter for the subtraction of \c RHS from \c LHS
523   Counter subtractCounters(Counter LHS, Counter RHS) {
524     return Builder.subtract(LHS, RHS);
525   }
526 
527   /// Return a counter for the sum of \c LHS and \c RHS.
528   Counter addCounters(Counter LHS, Counter RHS) {
529     return Builder.add(LHS, RHS);
530   }
531 
532   Counter addCounters(Counter C1, Counter C2, Counter C3) {
533     return addCounters(addCounters(C1, C2), C3);
534   }
535 
536   /// Return the region counter for the given statement.
537   ///
538   /// This should only be called on statements that have a dedicated counter.
539   Counter getRegionCounter(const Stmt *S) {
540     return Counter::getCounter(CounterMap[S]);
541   }
542 
543   /// Push a region onto the stack.
544   ///
545   /// Returns the index on the stack where the region was pushed. This can be
546   /// used with popRegions to exit a "scope", ending the region that was pushed.
547   size_t pushRegion(Counter Count, Optional<SourceLocation> StartLoc = None,
548                     Optional<SourceLocation> EndLoc = None) {
549     if (StartLoc) {
550       MostRecentLocation = *StartLoc;
551       completeDeferred(Count, MostRecentLocation);
552     }
553     RegionStack.emplace_back(Count, StartLoc, EndLoc);
554 
555     return RegionStack.size() - 1;
556   }
557 
558   /// Complete any pending deferred region by setting its end location and
559   /// count, and then pushing it onto the region stack.
560   size_t completeDeferred(Counter Count, SourceLocation DeferredEndLoc) {
561     size_t Index = RegionStack.size();
562     if (!DeferredRegion)
563       return Index;
564 
565     // Consume the pending region.
566     SourceMappingRegion DR = DeferredRegion.getValue();
567     DeferredRegion = None;
568 
569     // If the region ends in an expansion, find the expansion site.
570     FileID StartFile = SM.getFileID(DR.getBeginLoc());
571     if (SM.getFileID(DeferredEndLoc) != StartFile) {
572       if (isNestedIn(DeferredEndLoc, StartFile)) {
573         do {
574           DeferredEndLoc = getIncludeOrExpansionLoc(DeferredEndLoc);
575         } while (StartFile != SM.getFileID(DeferredEndLoc));
576       } else {
577         return Index;
578       }
579     }
580 
581     // The parent of this deferred region ends where the containing decl ends,
582     // so the region isn't useful.
583     if (DR.getBeginLoc() == DeferredEndLoc)
584       return Index;
585 
586     // If we're visiting statements in non-source order (e.g switch cases or
587     // a loop condition) we can't construct a sensible deferred region.
588     if (!SpellingRegion(SM, DR.getBeginLoc(), DeferredEndLoc).isInSourceOrder())
589       return Index;
590 
591     DR.setGap(true);
592     DR.setCounter(Count);
593     DR.setEndLoc(DeferredEndLoc);
594     handleFileExit(DeferredEndLoc);
595     RegionStack.push_back(DR);
596     return Index;
597   }
598 
599   /// Complete a deferred region created after a terminated region at the
600   /// top-level.
601   void completeTopLevelDeferredRegion(Counter Count,
602                                       SourceLocation DeferredEndLoc) {
603     if (DeferredRegion || !LastTerminatedRegion)
604       return;
605 
606     if (LastTerminatedRegion->second != RegionStack.size())
607       return;
608 
609     SourceLocation Start = LastTerminatedRegion->first;
610     if (SM.getFileID(Start) != SM.getMainFileID())
611       return;
612 
613     SourceMappingRegion DR = RegionStack.back();
614     DR.setStartLoc(Start);
615     DR.setDeferred(false);
616     DeferredRegion = DR;
617     completeDeferred(Count, DeferredEndLoc);
618   }
619 
620   size_t locationDepth(SourceLocation Loc) {
621     size_t Depth = 0;
622     while (Loc.isValid()) {
623       Loc = getIncludeOrExpansionLoc(Loc);
624       Depth++;
625     }
626     return Depth;
627   }
628 
629   /// Pop regions from the stack into the function's list of regions.
630   ///
631   /// Adds all regions from \c ParentIndex to the top of the stack to the
632   /// function's \c SourceRegions.
633   void popRegions(size_t ParentIndex) {
634     assert(RegionStack.size() >= ParentIndex && "parent not in stack");
635     bool ParentOfDeferredRegion = false;
636     while (RegionStack.size() > ParentIndex) {
637       SourceMappingRegion &Region = RegionStack.back();
638       if (Region.hasStartLoc()) {
639         SourceLocation StartLoc = Region.getBeginLoc();
640         SourceLocation EndLoc = Region.hasEndLoc()
641                                     ? Region.getEndLoc()
642                                     : RegionStack[ParentIndex].getEndLoc();
643         size_t StartDepth = locationDepth(StartLoc);
644         size_t EndDepth = locationDepth(EndLoc);
645         while (!SM.isWrittenInSameFile(StartLoc, EndLoc)) {
646           bool UnnestStart = StartDepth >= EndDepth;
647           bool UnnestEnd = EndDepth >= StartDepth;
648           if (UnnestEnd) {
649             // The region ends in a nested file or macro expansion. Create a
650             // separate region for each expansion.
651             SourceLocation NestedLoc = getStartOfFileOrMacro(EndLoc);
652             assert(SM.isWrittenInSameFile(NestedLoc, EndLoc));
653 
654             if (!isRegionAlreadyAdded(NestedLoc, EndLoc))
655               SourceRegions.emplace_back(Region.getCounter(), NestedLoc, EndLoc);
656 
657             EndLoc = getPreciseTokenLocEnd(getIncludeOrExpansionLoc(EndLoc));
658             if (EndLoc.isInvalid())
659               llvm::report_fatal_error("File exit not handled before popRegions");
660             EndDepth--;
661           }
662           if (UnnestStart) {
663             // The region begins in a nested file or macro expansion. Create a
664             // separate region for each expansion.
665             SourceLocation NestedLoc = getEndOfFileOrMacro(StartLoc);
666             assert(SM.isWrittenInSameFile(StartLoc, NestedLoc));
667 
668             if (!isRegionAlreadyAdded(StartLoc, NestedLoc))
669               SourceRegions.emplace_back(Region.getCounter(), StartLoc, NestedLoc);
670 
671             StartLoc = getIncludeOrExpansionLoc(StartLoc);
672             if (StartLoc.isInvalid())
673               llvm::report_fatal_error("File exit not handled before popRegions");
674             StartDepth--;
675           }
676         }
677         Region.setStartLoc(StartLoc);
678         Region.setEndLoc(EndLoc);
679 
680         MostRecentLocation = EndLoc;
681         // If this region happens to span an entire expansion, we need to make
682         // sure we don't overlap the parent region with it.
683         if (StartLoc == getStartOfFileOrMacro(StartLoc) &&
684             EndLoc == getEndOfFileOrMacro(EndLoc))
685           MostRecentLocation = getIncludeOrExpansionLoc(EndLoc);
686 
687         assert(SM.isWrittenInSameFile(Region.getBeginLoc(), EndLoc));
688         assert(SpellingRegion(SM, Region).isInSourceOrder());
689         SourceRegions.push_back(Region);
690 
691         if (ParentOfDeferredRegion) {
692           ParentOfDeferredRegion = false;
693 
694           // If there's an existing deferred region, keep the old one, because
695           // it means there are two consecutive returns (or a similar pattern).
696           if (!DeferredRegion.hasValue() &&
697               // File IDs aren't gathered within macro expansions, so it isn't
698               // useful to try and create a deferred region inside of one.
699               !EndLoc.isMacroID())
700             DeferredRegion =
701                 SourceMappingRegion(Counter::getZero(), EndLoc, None);
702         }
703       } else if (Region.isDeferred()) {
704         assert(!ParentOfDeferredRegion && "Consecutive deferred regions");
705         ParentOfDeferredRegion = true;
706       }
707       RegionStack.pop_back();
708 
709       // If the zero region pushed after the last terminated region no longer
710       // exists, clear its cached information.
711       if (LastTerminatedRegion &&
712           RegionStack.size() < LastTerminatedRegion->second)
713         LastTerminatedRegion = None;
714     }
715     assert(!ParentOfDeferredRegion && "Deferred region with no parent");
716   }
717 
718   /// Return the currently active region.
719   SourceMappingRegion &getRegion() {
720     assert(!RegionStack.empty() && "statement has no region");
721     return RegionStack.back();
722   }
723 
724   /// Propagate counts through the children of \p S if \p VisitChildren is true.
725   /// Otherwise, only emit a count for \p S itself.
726   Counter propagateCounts(Counter TopCount, const Stmt *S,
727                           bool VisitChildren = true) {
728     SourceLocation StartLoc = getStart(S);
729     SourceLocation EndLoc = getEnd(S);
730     size_t Index = pushRegion(TopCount, StartLoc, EndLoc);
731     if (VisitChildren)
732       Visit(S);
733     Counter ExitCount = getRegion().getCounter();
734     popRegions(Index);
735 
736     // The statement may be spanned by an expansion. Make sure we handle a file
737     // exit out of this expansion before moving to the next statement.
738     if (SM.isBeforeInTranslationUnit(StartLoc, S->getBeginLoc()))
739       MostRecentLocation = EndLoc;
740 
741     return ExitCount;
742   }
743 
744   /// Check whether a region with bounds \c StartLoc and \c EndLoc
745   /// is already added to \c SourceRegions.
746   bool isRegionAlreadyAdded(SourceLocation StartLoc, SourceLocation EndLoc) {
747     return SourceRegions.rend() !=
748            std::find_if(SourceRegions.rbegin(), SourceRegions.rend(),
749                         [&](const SourceMappingRegion &Region) {
750                           return Region.getBeginLoc() == StartLoc &&
751                                  Region.getEndLoc() == EndLoc;
752                         });
753   }
754 
755   /// Adjust the most recently visited location to \c EndLoc.
756   ///
757   /// This should be used after visiting any statements in non-source order.
758   void adjustForOutOfOrderTraversal(SourceLocation EndLoc) {
759     MostRecentLocation = EndLoc;
760     // The code region for a whole macro is created in handleFileExit() when
761     // it detects exiting of the virtual file of that macro. If we visited
762     // statements in non-source order, we might already have such a region
763     // added, for example, if a body of a loop is divided among multiple
764     // macros. Avoid adding duplicate regions in such case.
765     if (getRegion().hasEndLoc() &&
766         MostRecentLocation == getEndOfFileOrMacro(MostRecentLocation) &&
767         isRegionAlreadyAdded(getStartOfFileOrMacro(MostRecentLocation),
768                              MostRecentLocation))
769       MostRecentLocation = getIncludeOrExpansionLoc(MostRecentLocation);
770   }
771 
772   /// Adjust regions and state when \c NewLoc exits a file.
773   ///
774   /// If moving from our most recently tracked location to \c NewLoc exits any
775   /// files, this adjusts our current region stack and creates the file regions
776   /// for the exited file.
777   void handleFileExit(SourceLocation NewLoc) {
778     if (NewLoc.isInvalid() ||
779         SM.isWrittenInSameFile(MostRecentLocation, NewLoc))
780       return;
781 
782     // If NewLoc is not in a file that contains MostRecentLocation, walk up to
783     // find the common ancestor.
784     SourceLocation LCA = NewLoc;
785     FileID ParentFile = SM.getFileID(LCA);
786     while (!isNestedIn(MostRecentLocation, ParentFile)) {
787       LCA = getIncludeOrExpansionLoc(LCA);
788       if (LCA.isInvalid() || SM.isWrittenInSameFile(LCA, MostRecentLocation)) {
789         // Since there isn't a common ancestor, no file was exited. We just need
790         // to adjust our location to the new file.
791         MostRecentLocation = NewLoc;
792         return;
793       }
794       ParentFile = SM.getFileID(LCA);
795     }
796 
797     llvm::SmallSet<SourceLocation, 8> StartLocs;
798     Optional<Counter> ParentCounter;
799     for (SourceMappingRegion &I : llvm::reverse(RegionStack)) {
800       if (!I.hasStartLoc())
801         continue;
802       SourceLocation Loc = I.getBeginLoc();
803       if (!isNestedIn(Loc, ParentFile)) {
804         ParentCounter = I.getCounter();
805         break;
806       }
807 
808       while (!SM.isInFileID(Loc, ParentFile)) {
809         // The most nested region for each start location is the one with the
810         // correct count. We avoid creating redundant regions by stopping once
811         // we've seen this region.
812         if (StartLocs.insert(Loc).second)
813           SourceRegions.emplace_back(I.getCounter(), Loc,
814                                      getEndOfFileOrMacro(Loc));
815         Loc = getIncludeOrExpansionLoc(Loc);
816       }
817       I.setStartLoc(getPreciseTokenLocEnd(Loc));
818     }
819 
820     if (ParentCounter) {
821       // If the file is contained completely by another region and doesn't
822       // immediately start its own region, the whole file gets a region
823       // corresponding to the parent.
824       SourceLocation Loc = MostRecentLocation;
825       while (isNestedIn(Loc, ParentFile)) {
826         SourceLocation FileStart = getStartOfFileOrMacro(Loc);
827         if (StartLocs.insert(FileStart).second) {
828           SourceRegions.emplace_back(*ParentCounter, FileStart,
829                                      getEndOfFileOrMacro(Loc));
830           assert(SpellingRegion(SM, SourceRegions.back()).isInSourceOrder());
831         }
832         Loc = getIncludeOrExpansionLoc(Loc);
833       }
834     }
835 
836     MostRecentLocation = NewLoc;
837   }
838 
839   /// Ensure that \c S is included in the current region.
840   void extendRegion(const Stmt *S) {
841     SourceMappingRegion &Region = getRegion();
842     SourceLocation StartLoc = getStart(S);
843 
844     handleFileExit(StartLoc);
845     if (!Region.hasStartLoc())
846       Region.setStartLoc(StartLoc);
847 
848     completeDeferred(Region.getCounter(), StartLoc);
849   }
850 
851   /// Mark \c S as a terminator, starting a zero region.
852   void terminateRegion(const Stmt *S) {
853     extendRegion(S);
854     SourceMappingRegion &Region = getRegion();
855     SourceLocation EndLoc = getEnd(S);
856     if (!Region.hasEndLoc())
857       Region.setEndLoc(EndLoc);
858     pushRegion(Counter::getZero());
859     auto &ZeroRegion = getRegion();
860     ZeroRegion.setDeferred(true);
861     LastTerminatedRegion = {EndLoc, RegionStack.size()};
862   }
863 
864   /// Find a valid gap range between \p AfterLoc and \p BeforeLoc.
865   Optional<SourceRange> findGapAreaBetween(SourceLocation AfterLoc,
866                                            SourceLocation BeforeLoc) {
867     AfterLoc = SM.getExpansionLoc(AfterLoc);
868     BeforeLoc = SM.getExpansionLoc(BeforeLoc);
869     if (!SM.isWrittenInSameFile(AfterLoc, BeforeLoc))
870       return None;
871     return {{AfterLoc, BeforeLoc}};
872   }
873 
874   /// Emit a gap region between \p StartLoc and \p EndLoc with the given count.
875   void fillGapAreaWithCount(SourceLocation StartLoc, SourceLocation EndLoc,
876                             Counter Count) {
877     if (StartLoc == EndLoc)
878       return;
879     assert(SpellingRegion(SM, StartLoc, EndLoc).isInSourceOrder());
880     handleFileExit(StartLoc);
881     size_t Index = pushRegion(Count, StartLoc, EndLoc);
882     getRegion().setGap(true);
883     handleFileExit(EndLoc);
884     popRegions(Index);
885   }
886 
887   /// Keep counts of breaks and continues inside loops.
888   struct BreakContinue {
889     Counter BreakCount;
890     Counter ContinueCount;
891   };
892   SmallVector<BreakContinue, 8> BreakContinueStack;
893 
894   CounterCoverageMappingBuilder(
895       CoverageMappingModuleGen &CVM,
896       llvm::DenseMap<const Stmt *, unsigned> &CounterMap, SourceManager &SM,
897       const LangOptions &LangOpts)
898       : CoverageMappingBuilder(CVM, SM, LangOpts), CounterMap(CounterMap),
899         DeferredRegion(None) {}
900 
901   /// Write the mapping data to the output stream
902   void write(llvm::raw_ostream &OS) {
903     llvm::SmallVector<unsigned, 8> VirtualFileMapping;
904     gatherFileIDs(VirtualFileMapping);
905     SourceRegionFilter Filter = emitExpansionRegions();
906     assert(!DeferredRegion && "Deferred region never completed");
907     emitSourceRegions(Filter);
908     gatherSkippedRegions();
909 
910     if (MappingRegions.empty())
911       return;
912 
913     CoverageMappingWriter Writer(VirtualFileMapping, Builder.getExpressions(),
914                                  MappingRegions);
915     Writer.write(OS);
916   }
917 
918   void VisitStmt(const Stmt *S) {
919     if (S->getBeginLoc().isValid())
920       extendRegion(S);
921     for (const Stmt *Child : S->children())
922       if (Child)
923         this->Visit(Child);
924     handleFileExit(getEnd(S));
925   }
926 
927   void VisitDecl(const Decl *D) {
928     assert(!DeferredRegion && "Deferred region never completed");
929 
930     Stmt *Body = D->getBody();
931 
932     // Do not propagate region counts into system headers.
933     if (Body && SM.isInSystemHeader(SM.getSpellingLoc(getStart(Body))))
934       return;
935 
936     // Do not visit the artificial children nodes of defaulted methods. The
937     // lexer may not be able to report back precise token end locations for
938     // these children nodes (llvm.org/PR39822), and moreover users will not be
939     // able to see coverage for them.
940     bool Defaulted = false;
941     if (auto *Method = dyn_cast<CXXMethodDecl>(D))
942       Defaulted = Method->isDefaulted();
943 
944     propagateCounts(getRegionCounter(Body), Body,
945                     /*VisitChildren=*/!Defaulted);
946     assert(RegionStack.empty() && "Regions entered but never exited");
947 
948     // Discard the last uncompleted deferred region in a decl, if one exists.
949     // This prevents lines at the end of a function containing only whitespace
950     // or closing braces from being marked as uncovered.
951     DeferredRegion = None;
952   }
953 
954   void VisitReturnStmt(const ReturnStmt *S) {
955     extendRegion(S);
956     if (S->getRetValue())
957       Visit(S->getRetValue());
958     terminateRegion(S);
959   }
960 
961   void VisitCoroutineBodyStmt(const CoroutineBodyStmt *S) {
962     extendRegion(S);
963     Visit(S->getBody());
964   }
965 
966   void VisitCoreturnStmt(const CoreturnStmt *S) {
967     extendRegion(S);
968     if (S->getOperand())
969       Visit(S->getOperand());
970     terminateRegion(S);
971   }
972 
973   void VisitCXXThrowExpr(const CXXThrowExpr *E) {
974     extendRegion(E);
975     if (E->getSubExpr())
976       Visit(E->getSubExpr());
977     terminateRegion(E);
978   }
979 
980   void VisitGotoStmt(const GotoStmt *S) { terminateRegion(S); }
981 
982   void VisitLabelStmt(const LabelStmt *S) {
983     Counter LabelCount = getRegionCounter(S);
984     SourceLocation Start = getStart(S);
985     completeTopLevelDeferredRegion(LabelCount, Start);
986     completeDeferred(LabelCount, Start);
987     // We can't extendRegion here or we risk overlapping with our new region.
988     handleFileExit(Start);
989     pushRegion(LabelCount, Start);
990     Visit(S->getSubStmt());
991   }
992 
993   void VisitBreakStmt(const BreakStmt *S) {
994     assert(!BreakContinueStack.empty() && "break not in a loop or switch!");
995     BreakContinueStack.back().BreakCount = addCounters(
996         BreakContinueStack.back().BreakCount, getRegion().getCounter());
997     // FIXME: a break in a switch should terminate regions for all preceding
998     // case statements, not just the most recent one.
999     terminateRegion(S);
1000   }
1001 
1002   void VisitContinueStmt(const ContinueStmt *S) {
1003     assert(!BreakContinueStack.empty() && "continue stmt not in a loop!");
1004     BreakContinueStack.back().ContinueCount = addCounters(
1005         BreakContinueStack.back().ContinueCount, getRegion().getCounter());
1006     terminateRegion(S);
1007   }
1008 
1009   void VisitCallExpr(const CallExpr *E) {
1010     VisitStmt(E);
1011 
1012     // Terminate the region when we hit a noreturn function.
1013     // (This is helpful dealing with switch statements.)
1014     QualType CalleeType = E->getCallee()->getType();
1015     if (getFunctionExtInfo(*CalleeType).getNoReturn())
1016       terminateRegion(E);
1017   }
1018 
1019   void VisitWhileStmt(const WhileStmt *S) {
1020     extendRegion(S);
1021 
1022     Counter ParentCount = getRegion().getCounter();
1023     Counter BodyCount = getRegionCounter(S);
1024 
1025     // Handle the body first so that we can get the backedge count.
1026     BreakContinueStack.push_back(BreakContinue());
1027     extendRegion(S->getBody());
1028     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
1029     BreakContinue BC = BreakContinueStack.pop_back_val();
1030 
1031     // Go back to handle the condition.
1032     Counter CondCount =
1033         addCounters(ParentCount, BackedgeCount, BC.ContinueCount);
1034     propagateCounts(CondCount, S->getCond());
1035     adjustForOutOfOrderTraversal(getEnd(S));
1036 
1037     // The body count applies to the area immediately after the increment.
1038     auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
1039                                   getStart(S->getBody()));
1040     if (Gap)
1041       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
1042 
1043     Counter OutCount =
1044         addCounters(BC.BreakCount, subtractCounters(CondCount, BodyCount));
1045     if (OutCount != ParentCount)
1046       pushRegion(OutCount);
1047   }
1048 
1049   void VisitDoStmt(const DoStmt *S) {
1050     extendRegion(S);
1051 
1052     Counter ParentCount = getRegion().getCounter();
1053     Counter BodyCount = getRegionCounter(S);
1054 
1055     BreakContinueStack.push_back(BreakContinue());
1056     extendRegion(S->getBody());
1057     Counter BackedgeCount =
1058         propagateCounts(addCounters(ParentCount, BodyCount), S->getBody());
1059     BreakContinue BC = BreakContinueStack.pop_back_val();
1060 
1061     Counter CondCount = addCounters(BackedgeCount, BC.ContinueCount);
1062     propagateCounts(CondCount, S->getCond());
1063 
1064     Counter OutCount =
1065         addCounters(BC.BreakCount, subtractCounters(CondCount, BodyCount));
1066     if (OutCount != ParentCount)
1067       pushRegion(OutCount);
1068   }
1069 
1070   void VisitForStmt(const ForStmt *S) {
1071     extendRegion(S);
1072     if (S->getInit())
1073       Visit(S->getInit());
1074 
1075     Counter ParentCount = getRegion().getCounter();
1076     Counter BodyCount = getRegionCounter(S);
1077 
1078     // The loop increment may contain a break or continue.
1079     if (S->getInc())
1080       BreakContinueStack.emplace_back();
1081 
1082     // Handle the body first so that we can get the backedge count.
1083     BreakContinueStack.emplace_back();
1084     extendRegion(S->getBody());
1085     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
1086     BreakContinue BodyBC = BreakContinueStack.pop_back_val();
1087 
1088     // The increment is essentially part of the body but it needs to include
1089     // the count for all the continue statements.
1090     BreakContinue IncrementBC;
1091     if (const Stmt *Inc = S->getInc()) {
1092       propagateCounts(addCounters(BackedgeCount, BodyBC.ContinueCount), Inc);
1093       IncrementBC = BreakContinueStack.pop_back_val();
1094     }
1095 
1096     // Go back to handle the condition.
1097     Counter CondCount = addCounters(
1098         addCounters(ParentCount, BackedgeCount, BodyBC.ContinueCount),
1099         IncrementBC.ContinueCount);
1100     if (const Expr *Cond = S->getCond()) {
1101       propagateCounts(CondCount, Cond);
1102       adjustForOutOfOrderTraversal(getEnd(S));
1103     }
1104 
1105     // The body count applies to the area immediately after the increment.
1106     auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
1107                                   getStart(S->getBody()));
1108     if (Gap)
1109       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
1110 
1111     Counter OutCount = addCounters(BodyBC.BreakCount, IncrementBC.BreakCount,
1112                                    subtractCounters(CondCount, BodyCount));
1113     if (OutCount != ParentCount)
1114       pushRegion(OutCount);
1115   }
1116 
1117   void VisitCXXForRangeStmt(const CXXForRangeStmt *S) {
1118     extendRegion(S);
1119     if (S->getInit())
1120       Visit(S->getInit());
1121     Visit(S->getLoopVarStmt());
1122     Visit(S->getRangeStmt());
1123 
1124     Counter ParentCount = getRegion().getCounter();
1125     Counter BodyCount = getRegionCounter(S);
1126 
1127     BreakContinueStack.push_back(BreakContinue());
1128     extendRegion(S->getBody());
1129     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
1130     BreakContinue BC = BreakContinueStack.pop_back_val();
1131 
1132     // The body count applies to the area immediately after the range.
1133     auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
1134                                   getStart(S->getBody()));
1135     if (Gap)
1136       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
1137 
1138     Counter LoopCount =
1139         addCounters(ParentCount, BackedgeCount, BC.ContinueCount);
1140     Counter OutCount =
1141         addCounters(BC.BreakCount, subtractCounters(LoopCount, BodyCount));
1142     if (OutCount != ParentCount)
1143       pushRegion(OutCount);
1144   }
1145 
1146   void VisitObjCForCollectionStmt(const ObjCForCollectionStmt *S) {
1147     extendRegion(S);
1148     Visit(S->getElement());
1149 
1150     Counter ParentCount = getRegion().getCounter();
1151     Counter BodyCount = getRegionCounter(S);
1152 
1153     BreakContinueStack.push_back(BreakContinue());
1154     extendRegion(S->getBody());
1155     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
1156     BreakContinue BC = BreakContinueStack.pop_back_val();
1157 
1158     // The body count applies to the area immediately after the collection.
1159     auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
1160                                   getStart(S->getBody()));
1161     if (Gap)
1162       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
1163 
1164     Counter LoopCount =
1165         addCounters(ParentCount, BackedgeCount, BC.ContinueCount);
1166     Counter OutCount =
1167         addCounters(BC.BreakCount, subtractCounters(LoopCount, BodyCount));
1168     if (OutCount != ParentCount)
1169       pushRegion(OutCount);
1170   }
1171 
1172   void VisitSwitchStmt(const SwitchStmt *S) {
1173     extendRegion(S);
1174     if (S->getInit())
1175       Visit(S->getInit());
1176     Visit(S->getCond());
1177 
1178     BreakContinueStack.push_back(BreakContinue());
1179 
1180     const Stmt *Body = S->getBody();
1181     extendRegion(Body);
1182     if (const auto *CS = dyn_cast<CompoundStmt>(Body)) {
1183       if (!CS->body_empty()) {
1184         // Make a region for the body of the switch.  If the body starts with
1185         // a case, that case will reuse this region; otherwise, this covers
1186         // the unreachable code at the beginning of the switch body.
1187         size_t Index = pushRegion(Counter::getZero(), getStart(CS));
1188         getRegion().setGap(true);
1189         for (const auto *Child : CS->children())
1190           Visit(Child);
1191 
1192         // Set the end for the body of the switch, if it isn't already set.
1193         for (size_t i = RegionStack.size(); i != Index; --i) {
1194           if (!RegionStack[i - 1].hasEndLoc())
1195             RegionStack[i - 1].setEndLoc(getEnd(CS->body_back()));
1196         }
1197 
1198         popRegions(Index);
1199       }
1200     } else
1201       propagateCounts(Counter::getZero(), Body);
1202     BreakContinue BC = BreakContinueStack.pop_back_val();
1203 
1204     if (!BreakContinueStack.empty())
1205       BreakContinueStack.back().ContinueCount = addCounters(
1206           BreakContinueStack.back().ContinueCount, BC.ContinueCount);
1207 
1208     Counter ExitCount = getRegionCounter(S);
1209     SourceLocation ExitLoc = getEnd(S);
1210     pushRegion(ExitCount);
1211 
1212     // Ensure that handleFileExit recognizes when the end location is located
1213     // in a different file.
1214     MostRecentLocation = getStart(S);
1215     handleFileExit(ExitLoc);
1216   }
1217 
1218   void VisitSwitchCase(const SwitchCase *S) {
1219     extendRegion(S);
1220 
1221     SourceMappingRegion &Parent = getRegion();
1222 
1223     Counter Count = addCounters(Parent.getCounter(), getRegionCounter(S));
1224     // Reuse the existing region if it starts at our label. This is typical of
1225     // the first case in a switch.
1226     if (Parent.hasStartLoc() && Parent.getBeginLoc() == getStart(S))
1227       Parent.setCounter(Count);
1228     else
1229       pushRegion(Count, getStart(S));
1230 
1231     if (const auto *CS = dyn_cast<CaseStmt>(S)) {
1232       Visit(CS->getLHS());
1233       if (const Expr *RHS = CS->getRHS())
1234         Visit(RHS);
1235     }
1236     Visit(S->getSubStmt());
1237   }
1238 
1239   void VisitIfStmt(const IfStmt *S) {
1240     extendRegion(S);
1241     if (S->getInit())
1242       Visit(S->getInit());
1243 
1244     // Extend into the condition before we propagate through it below - this is
1245     // needed to handle macros that generate the "if" but not the condition.
1246     extendRegion(S->getCond());
1247 
1248     Counter ParentCount = getRegion().getCounter();
1249     Counter ThenCount = getRegionCounter(S);
1250 
1251     // Emitting a counter for the condition makes it easier to interpret the
1252     // counter for the body when looking at the coverage.
1253     propagateCounts(ParentCount, S->getCond());
1254 
1255     // The 'then' count applies to the area immediately after the condition.
1256     auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
1257                                   getStart(S->getThen()));
1258     if (Gap)
1259       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), ThenCount);
1260 
1261     extendRegion(S->getThen());
1262     Counter OutCount = propagateCounts(ThenCount, S->getThen());
1263 
1264     Counter ElseCount = subtractCounters(ParentCount, ThenCount);
1265     if (const Stmt *Else = S->getElse()) {
1266       // The 'else' count applies to the area immediately after the 'then'.
1267       Gap = findGapAreaBetween(getPreciseTokenLocEnd(getEnd(S->getThen())),
1268                                getStart(Else));
1269       if (Gap)
1270         fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), ElseCount);
1271       extendRegion(Else);
1272       OutCount = addCounters(OutCount, propagateCounts(ElseCount, Else));
1273     } else
1274       OutCount = addCounters(OutCount, ElseCount);
1275 
1276     if (OutCount != ParentCount)
1277       pushRegion(OutCount);
1278   }
1279 
1280   void VisitCXXTryStmt(const CXXTryStmt *S) {
1281     extendRegion(S);
1282     // Handle macros that generate the "try" but not the rest.
1283     extendRegion(S->getTryBlock());
1284 
1285     Counter ParentCount = getRegion().getCounter();
1286     propagateCounts(ParentCount, S->getTryBlock());
1287 
1288     for (unsigned I = 0, E = S->getNumHandlers(); I < E; ++I)
1289       Visit(S->getHandler(I));
1290 
1291     Counter ExitCount = getRegionCounter(S);
1292     pushRegion(ExitCount);
1293   }
1294 
1295   void VisitCXXCatchStmt(const CXXCatchStmt *S) {
1296     propagateCounts(getRegionCounter(S), S->getHandlerBlock());
1297   }
1298 
1299   void VisitAbstractConditionalOperator(const AbstractConditionalOperator *E) {
1300     extendRegion(E);
1301 
1302     Counter ParentCount = getRegion().getCounter();
1303     Counter TrueCount = getRegionCounter(E);
1304 
1305     Visit(E->getCond());
1306 
1307     if (!isa<BinaryConditionalOperator>(E)) {
1308       // The 'then' count applies to the area immediately after the condition.
1309       auto Gap =
1310           findGapAreaBetween(E->getQuestionLoc(), getStart(E->getTrueExpr()));
1311       if (Gap)
1312         fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), TrueCount);
1313 
1314       extendRegion(E->getTrueExpr());
1315       propagateCounts(TrueCount, E->getTrueExpr());
1316     }
1317 
1318     extendRegion(E->getFalseExpr());
1319     propagateCounts(subtractCounters(ParentCount, TrueCount),
1320                     E->getFalseExpr());
1321   }
1322 
1323   void VisitBinLAnd(const BinaryOperator *E) {
1324     extendRegion(E->getLHS());
1325     propagateCounts(getRegion().getCounter(), E->getLHS());
1326     handleFileExit(getEnd(E->getLHS()));
1327 
1328     extendRegion(E->getRHS());
1329     propagateCounts(getRegionCounter(E), E->getRHS());
1330   }
1331 
1332   void VisitBinLOr(const BinaryOperator *E) {
1333     extendRegion(E->getLHS());
1334     propagateCounts(getRegion().getCounter(), E->getLHS());
1335     handleFileExit(getEnd(E->getLHS()));
1336 
1337     extendRegion(E->getRHS());
1338     propagateCounts(getRegionCounter(E), E->getRHS());
1339   }
1340 
1341   void VisitLambdaExpr(const LambdaExpr *LE) {
1342     // Lambdas are treated as their own functions for now, so we shouldn't
1343     // propagate counts into them.
1344   }
1345 };
1346 
1347 std::string normalizeFilename(StringRef Filename) {
1348   llvm::SmallString<256> Path(Filename);
1349   llvm::sys::fs::make_absolute(Path);
1350   llvm::sys::path::remove_dots(Path, /*remove_dot_dot=*/true);
1351   return std::string(Path);
1352 }
1353 
1354 } // end anonymous namespace
1355 
1356 static void dump(llvm::raw_ostream &OS, StringRef FunctionName,
1357                  ArrayRef<CounterExpression> Expressions,
1358                  ArrayRef<CounterMappingRegion> Regions) {
1359   OS << FunctionName << ":\n";
1360   CounterMappingContext Ctx(Expressions);
1361   for (const auto &R : Regions) {
1362     OS.indent(2);
1363     switch (R.Kind) {
1364     case CounterMappingRegion::CodeRegion:
1365       break;
1366     case CounterMappingRegion::ExpansionRegion:
1367       OS << "Expansion,";
1368       break;
1369     case CounterMappingRegion::SkippedRegion:
1370       OS << "Skipped,";
1371       break;
1372     case CounterMappingRegion::GapRegion:
1373       OS << "Gap,";
1374       break;
1375     }
1376 
1377     OS << "File " << R.FileID << ", " << R.LineStart << ":" << R.ColumnStart
1378        << " -> " << R.LineEnd << ":" << R.ColumnEnd << " = ";
1379     Ctx.dump(R.Count, OS);
1380     if (R.Kind == CounterMappingRegion::ExpansionRegion)
1381       OS << " (Expanded file = " << R.ExpandedFileID << ")";
1382     OS << "\n";
1383   }
1384 }
1385 
1386 static std::string getInstrProfSection(const CodeGenModule &CGM,
1387                                        llvm::InstrProfSectKind SK) {
1388   return llvm::getInstrProfSectionName(
1389       SK, CGM.getContext().getTargetInfo().getTriple().getObjectFormat());
1390 }
1391 
1392 void CoverageMappingModuleGen::emitFunctionMappingRecord(
1393     const FunctionInfo &Info, uint64_t FilenamesRef) {
1394   llvm::LLVMContext &Ctx = CGM.getLLVMContext();
1395 
1396   // Assign a name to the function record. This is used to merge duplicates.
1397   std::string FuncRecordName = "__covrec_" + llvm::utohexstr(Info.NameHash);
1398 
1399   // A dummy description for a function included-but-not-used in a TU can be
1400   // replaced by full description provided by a different TU. The two kinds of
1401   // descriptions play distinct roles: therefore, assign them different names
1402   // to prevent `linkonce_odr` merging.
1403   if (Info.IsUsed)
1404     FuncRecordName += "u";
1405 
1406   // Create the function record type.
1407   const uint64_t NameHash = Info.NameHash;
1408   const uint64_t FuncHash = Info.FuncHash;
1409   const std::string &CoverageMapping = Info.CoverageMapping;
1410 #define COVMAP_FUNC_RECORD(Type, LLVMType, Name, Init) LLVMType,
1411   llvm::Type *FunctionRecordTypes[] = {
1412 #include "llvm/ProfileData/InstrProfData.inc"
1413   };
1414   auto *FunctionRecordTy =
1415       llvm::StructType::get(Ctx, makeArrayRef(FunctionRecordTypes),
1416                             /*isPacked=*/true);
1417 
1418   // Create the function record constant.
1419 #define COVMAP_FUNC_RECORD(Type, LLVMType, Name, Init) Init,
1420   llvm::Constant *FunctionRecordVals[] = {
1421       #include "llvm/ProfileData/InstrProfData.inc"
1422   };
1423   auto *FuncRecordConstant = llvm::ConstantStruct::get(
1424       FunctionRecordTy, makeArrayRef(FunctionRecordVals));
1425 
1426   // Create the function record global.
1427   auto *FuncRecord = new llvm::GlobalVariable(
1428       CGM.getModule(), FunctionRecordTy, /*isConstant=*/true,
1429       llvm::GlobalValue::LinkOnceODRLinkage, FuncRecordConstant,
1430       FuncRecordName);
1431   FuncRecord->setVisibility(llvm::GlobalValue::HiddenVisibility);
1432   FuncRecord->setSection(getInstrProfSection(CGM, llvm::IPSK_covfun));
1433   FuncRecord->setAlignment(llvm::Align(8));
1434   if (CGM.supportsCOMDAT())
1435     FuncRecord->setComdat(CGM.getModule().getOrInsertComdat(FuncRecordName));
1436 
1437   // Make sure the data doesn't get deleted.
1438   CGM.addUsedGlobal(FuncRecord);
1439 }
1440 
1441 void CoverageMappingModuleGen::addFunctionMappingRecord(
1442     llvm::GlobalVariable *NamePtr, StringRef NameValue, uint64_t FuncHash,
1443     const std::string &CoverageMapping, bool IsUsed) {
1444   llvm::LLVMContext &Ctx = CGM.getLLVMContext();
1445   const uint64_t NameHash = llvm::IndexedInstrProf::ComputeHash(NameValue);
1446   FunctionRecords.push_back({NameHash, FuncHash, CoverageMapping, IsUsed});
1447 
1448   if (!IsUsed)
1449     FunctionNames.push_back(
1450         llvm::ConstantExpr::getBitCast(NamePtr, llvm::Type::getInt8PtrTy(Ctx)));
1451 
1452   if (CGM.getCodeGenOpts().DumpCoverageMapping) {
1453     // Dump the coverage mapping data for this function by decoding the
1454     // encoded data. This allows us to dump the mapping regions which were
1455     // also processed by the CoverageMappingWriter which performs
1456     // additional minimization operations such as reducing the number of
1457     // expressions.
1458     std::vector<StringRef> Filenames;
1459     std::vector<CounterExpression> Expressions;
1460     std::vector<CounterMappingRegion> Regions;
1461     llvm::SmallVector<std::string, 16> FilenameStrs;
1462     llvm::SmallVector<StringRef, 16> FilenameRefs;
1463     FilenameStrs.resize(FileEntries.size());
1464     FilenameRefs.resize(FileEntries.size());
1465     for (const auto &Entry : FileEntries) {
1466       auto I = Entry.second;
1467       FilenameStrs[I] = normalizeFilename(Entry.first->getName());
1468       FilenameRefs[I] = FilenameStrs[I];
1469     }
1470     RawCoverageMappingReader Reader(CoverageMapping, FilenameRefs, Filenames,
1471                                     Expressions, Regions);
1472     if (Reader.read())
1473       return;
1474     dump(llvm::outs(), NameValue, Expressions, Regions);
1475   }
1476 }
1477 
1478 void CoverageMappingModuleGen::emit() {
1479   if (FunctionRecords.empty())
1480     return;
1481   llvm::LLVMContext &Ctx = CGM.getLLVMContext();
1482   auto *Int32Ty = llvm::Type::getInt32Ty(Ctx);
1483 
1484   // Create the filenames and merge them with coverage mappings
1485   llvm::SmallVector<std::string, 16> FilenameStrs;
1486   llvm::SmallVector<StringRef, 16> FilenameRefs;
1487   FilenameStrs.resize(FileEntries.size());
1488   FilenameRefs.resize(FileEntries.size());
1489   for (const auto &Entry : FileEntries) {
1490     auto I = Entry.second;
1491     FilenameStrs[I] = normalizeFilename(Entry.first->getName());
1492     FilenameRefs[I] = FilenameStrs[I];
1493   }
1494 
1495   std::string Filenames;
1496   {
1497     llvm::raw_string_ostream OS(Filenames);
1498     CoverageFilenamesSectionWriter(FilenameRefs).write(OS);
1499   }
1500   auto *FilenamesVal =
1501       llvm::ConstantDataArray::getString(Ctx, Filenames, false);
1502   const int64_t FilenamesRef = llvm::IndexedInstrProf::ComputeHash(Filenames);
1503 
1504   // Emit the function records.
1505   for (const FunctionInfo &Info : FunctionRecords)
1506     emitFunctionMappingRecord(Info, FilenamesRef);
1507 
1508   const unsigned NRecords = 0;
1509   const size_t FilenamesSize = Filenames.size();
1510   const unsigned CoverageMappingSize = 0;
1511   llvm::Type *CovDataHeaderTypes[] = {
1512 #define COVMAP_HEADER(Type, LLVMType, Name, Init) LLVMType,
1513 #include "llvm/ProfileData/InstrProfData.inc"
1514   };
1515   auto CovDataHeaderTy =
1516       llvm::StructType::get(Ctx, makeArrayRef(CovDataHeaderTypes));
1517   llvm::Constant *CovDataHeaderVals[] = {
1518 #define COVMAP_HEADER(Type, LLVMType, Name, Init) Init,
1519 #include "llvm/ProfileData/InstrProfData.inc"
1520   };
1521   auto CovDataHeaderVal = llvm::ConstantStruct::get(
1522       CovDataHeaderTy, makeArrayRef(CovDataHeaderVals));
1523 
1524   // Create the coverage data record
1525   llvm::Type *CovDataTypes[] = {CovDataHeaderTy, FilenamesVal->getType()};
1526   auto CovDataTy = llvm::StructType::get(Ctx, makeArrayRef(CovDataTypes));
1527   llvm::Constant *TUDataVals[] = {CovDataHeaderVal, FilenamesVal};
1528   auto CovDataVal =
1529       llvm::ConstantStruct::get(CovDataTy, makeArrayRef(TUDataVals));
1530   auto CovData = new llvm::GlobalVariable(
1531       CGM.getModule(), CovDataTy, true, llvm::GlobalValue::PrivateLinkage,
1532       CovDataVal, llvm::getCoverageMappingVarName());
1533 
1534   CovData->setSection(getInstrProfSection(CGM, llvm::IPSK_covmap));
1535   CovData->setAlignment(llvm::Align(8));
1536 
1537   // Make sure the data doesn't get deleted.
1538   CGM.addUsedGlobal(CovData);
1539   // Create the deferred function records array
1540   if (!FunctionNames.empty()) {
1541     auto NamesArrTy = llvm::ArrayType::get(llvm::Type::getInt8PtrTy(Ctx),
1542                                            FunctionNames.size());
1543     auto NamesArrVal = llvm::ConstantArray::get(NamesArrTy, FunctionNames);
1544     // This variable will *NOT* be emitted to the object file. It is used
1545     // to pass the list of names referenced to codegen.
1546     new llvm::GlobalVariable(CGM.getModule(), NamesArrTy, true,
1547                              llvm::GlobalValue::InternalLinkage, NamesArrVal,
1548                              llvm::getCoverageUnusedNamesVarName());
1549   }
1550 }
1551 
1552 unsigned CoverageMappingModuleGen::getFileID(const FileEntry *File) {
1553   auto It = FileEntries.find(File);
1554   if (It != FileEntries.end())
1555     return It->second;
1556   unsigned FileID = FileEntries.size();
1557   FileEntries.insert(std::make_pair(File, FileID));
1558   return FileID;
1559 }
1560 
1561 void CoverageMappingGen::emitCounterMapping(const Decl *D,
1562                                             llvm::raw_ostream &OS) {
1563   assert(CounterMap);
1564   CounterCoverageMappingBuilder Walker(CVM, *CounterMap, SM, LangOpts);
1565   Walker.VisitDecl(D);
1566   Walker.write(OS);
1567 }
1568 
1569 void CoverageMappingGen::emitEmptyMapping(const Decl *D,
1570                                           llvm::raw_ostream &OS) {
1571   EmptyCoverageMappingBuilder Walker(CVM, SM, LangOpts);
1572   Walker.VisitDecl(D);
1573   Walker.write(OS);
1574 }
1575