[Profiler] Track the parent SourceFile for a SILCoverageMap

This tracks the source file that the profiled
entity was defined in.

Author: hamishknight

include/swift/SIL/SILCoverageMap.h

@@ -55,6 +55,15 @@ class SILCoverageMap : public llvm::ilist_node<SILCoverageMap>,
   };
 
 private:
+  /// The parent source file containing the coverage map.
+  /// 
+  /// NOTE: `ParentSourceFile->getFilename()` is not necessarily equivalent to
+  /// `Filename`. `Filename` could be a .swift file, and `ParentSourceFile`
+  /// could be a parsed .sil file. As such, this should only be used for
+  /// determining which object file to emit the coverage map into.
+  /// `Filename` should be used for coverage data.
+  SourceFile *ParentSourceFile;
+
   // The name of the source file where this mapping is found.
   StringRef Filename;
 
@@ -78,19 +87,28 @@ class SILCoverageMap : public llvm::ilist_node<SILCoverageMap>,
   SILCoverageMap &operator=(const SILCoverageMap &) = delete;
 
   /// Private constructor. Create these using SILCoverageMap::create.
-  SILCoverageMap(uint64_t Hash);
+  SILCoverageMap(SourceFile *ParentSourceFile, uint64_t Hash);
 
 public:
   ~SILCoverageMap();
 
   static SILCoverageMap *
-  create(SILModule &M, StringRef Filename, StringRef Name,
-         StringRef PGOFuncName, uint64_t Hash,
+  create(SILModule &M, SourceFile *ParentSourceFile, StringRef Filename,
+         StringRef Name, StringRef PGOFuncName, uint64_t Hash,
          ArrayRef<MappedRegion> MappedRegions,
          ArrayRef<llvm::coverage::CounterExpression> Expressions);
 
+  /// The parent source file containing the coverage map.
+  ///
+  /// NOTE: `getParentSourceFile()->getFilename()` is not necessarily equivalent
+  /// to `getFilename()`. `getFilename()` could be a .swift file, and
+  /// `getParentSourceFile()` could be a parsed .sil file. As such, this should
+  /// only be used for determining which object file to emit the coverage map
+  /// in. `getFilename()` should be used for coverage data.
+  SourceFile *getParentSourceFile() const { return ParentSourceFile; }
+
   /// Return the name of the source file where this mapping is found.
-  StringRef getFile() const { return Filename; }
+  StringRef getFilename() const { return Filename; }
 
   /// Return the mangled name of the function this mapping covers.
   StringRef getName() const { return Name; }

lib/IRGen/GenCoverage.cpp

@@ -80,8 +80,8 @@ void IRGenModule::emitCoverageMapping() {
 
   std::vector<StringRef> Files;
   for (const auto &M : Mappings)
-    if (std::find(Files.begin(), Files.end(), M->getFile()) == Files.end())
-      Files.push_back(M->getFile());
+    if (std::find(Files.begin(), Files.end(), M->getFilename()) == Files.end())
+      Files.push_back(M->getFilename());
 
   auto remapper = getOptions().CoveragePrefixMap;
 
@@ -115,7 +115,7 @@ void IRGenModule::emitCoverageMapping() {
     std::string FuncRecordName = "__covrec_" + llvm::utohexstr(NameHash);
 
     unsigned FileID =
-        std::find(Files.begin(), Files.end(), M->getFile()) - Files.begin();
+        std::find(Files.begin(), Files.end(), M->getFilename()) - Files.begin();
     std::vector<CounterMappingRegion> Regions;
     for (const auto &MR : M->getMappedRegions())
       Regions.emplace_back(CounterMappingRegion::makeRegion(

lib/SIL/IR/SILCoverageMap.cpp

@@ -24,12 +24,13 @@ using namespace swift;
 using llvm::coverage::CounterExpression;
 
 SILCoverageMap *
-SILCoverageMap::create(SILModule &M, StringRef Filename, StringRef Name,
+SILCoverageMap::create(SILModule &M, SourceFile *ParentSourceFile,
+                       StringRef Filename, StringRef Name,
                        StringRef PGOFuncName, uint64_t Hash,
                        ArrayRef<MappedRegion> MappedRegions,
                        ArrayRef<CounterExpression> Expressions) {
   auto *Buf = M.allocate<SILCoverageMap>(1);
-  SILCoverageMap *CM = ::new (Buf) SILCoverageMap(Hash);
+  SILCoverageMap *CM = ::new (Buf) SILCoverageMap(ParentSourceFile, Hash);
 
   // Store a copy of the names so that we own the lifetime.
   CM->Filename = M.allocateCopy(Filename);
@@ -50,7 +51,8 @@ SILCoverageMap::create(SILModule &M, StringRef Filename, StringRef Name,
   return CM;
 }
 
-SILCoverageMap::SILCoverageMap(uint64_t Hash) : Hash(Hash) {}
+SILCoverageMap::SILCoverageMap(SourceFile *ParentSourceFile, uint64_t Hash)
+  : ParentSourceFile(ParentSourceFile), Hash(Hash) {}
 
 SILCoverageMap::~SILCoverageMap() {}
 

lib/SIL/IR/SILPrinter.cpp

@@ -3877,7 +3877,7 @@ void SILDifferentiabilityWitness::dump() const {
 
 void SILCoverageMap::print(SILPrintContext &PrintCtx) const {
   llvm::raw_ostream &OS = PrintCtx.OS();
-  OS << "sil_coverage_map " << QuotedString(getFile()) << " "
+  OS << "sil_coverage_map " << QuotedString(getFilename()) << " "
      << QuotedString(getName()) << " " << QuotedString(getPGOFuncName()) << " "
      << getHash() << " {\t// " << demangleSymbol(getName()) << "\n";
   if (PrintCtx.sortSIL())

lib/SIL/IR/SILProfiler.cpp

@@ -1033,7 +1033,7 @@ struct CoverageMapping : public ASTWalker {
   SILCoverageMap *emitSourceRegions(
       SILModule &M, StringRef Name, StringRef PGOFuncName, uint64_t Hash,
       llvm::DenseMap<ProfileCounterRef, unsigned> &CounterIndices,
-      StringRef Filename) {
+      SourceFile *SF, StringRef Filename) {
     if (SourceRegions.empty())
       return nullptr;
 
@@ -1051,8 +1051,8 @@ struct CoverageMapping : public ASTWalker {
       Regions.emplace_back(Start.first, Start.second, End.first, End.second,
                            Counter.expand(Builder, CounterIndices));
     }
-    return SILCoverageMap::create(M, Filename, Name, PGOFuncName, Hash, Regions,
-                                  Builder.getExpressions());
+    return SILCoverageMap::create(M, SF, Filename, Name, PGOFuncName, Hash,
+                                  Regions, Builder.getExpressions());
   }
 
   PreWalkAction walkToDeclPre(Decl *D) override {
@@ -1331,13 +1331,6 @@ getEquivalentPGOLinkage(FormalLinkage Linkage) {
   llvm_unreachable("Unhandled FormalLinkage in switch.");
 }
 
-static StringRef getCurrentFileName(SILDeclRef forDecl) {
-  auto *DC = forDecl.getInnermostDeclContext();
-  if (auto *ParentFile = DC->getParentSourceFile())
-    return ParentFile->getFilename();
-  return {};
-}
-
 static void walkNode(NodeToProfile Node, ASTWalker &Walker) {
   if (auto N = Node.getAsNode()) {
     N.walk(Walker);
@@ -1350,8 +1343,11 @@ static void walkNode(NodeToProfile Node, ASTWalker &Walker) {
 
 void SILProfiler::assignRegionCounters() {
   const auto &SM = M.getASTContext().SourceMgr;
+  auto *DC = forDecl.getInnermostDeclContext();
+  auto *SF = DC->getParentSourceFile();
+  assert(SF && "Not within a SourceFile?");
 
-  CurrentFileName = getCurrentFileName(forDecl);
+  CurrentFileName = SF->getFilename();
 
   MapRegionCounters Mapper(forDecl, RegionCounterMap);
 
@@ -1387,7 +1383,7 @@ void SILProfiler::assignRegionCounters() {
     walkNode(Root, Coverage);
     CovMap =
         Coverage.emitSourceRegions(M, CurrentFuncName, PGOFuncName, PGOFuncHash,
-                                   RegionCounterMap, CurrentFileName);
+                                   RegionCounterMap, SF, CurrentFileName);
   }
 
   if (llvm::IndexedInstrProfReader *IPR = M.getPGOReader()) {

lib/SIL/Parser/ParseSIL.cpp

@@ -7946,8 +7946,9 @@ bool SILParserState::parseSILCoverageMap(Parser &P) {
                        LBraceLoc);
 
   if (!BodyHasError)
-    SILCoverageMap::create(M, Filename.str(), FuncName.str(), PGOFuncName.str(),
-                           Hash, Regions, Builder.getExpressions());
+    SILCoverageMap::create(M, &P.SF, Filename.str(), FuncName.str(),
+                           PGOFuncName.str(), Hash, Regions,
+                           Builder.getExpressions());
   return false;
 }