Merge pull request #16046 from vedantk/cov-visit-extensions

Author: swift-ci

lib/SIL/SILProfiler.cpp

@@ -93,11 +93,8 @@ static bool canCreateProfilerForAST(ASTNode N) {
   assert(hasASTBeenTypeChecked(N) && "Cannot use this AST for profiling");
 
   if (auto *D = N.dyn_cast<Decl *>()) {
-    // Any mapped function may be profiled. There's an exception for
-    // constructors because all of the constructors for a type share a single
-    // profiler.
     if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D))
-      return !isa<ConstructorDecl>(AFD);
+      return true;
 
     if (isa<TopLevelCodeDecl>(D))
       return true;
@@ -798,7 +795,7 @@ struct CoverageMapping : public ASTWalker {
       return visitFunctionDecl(*this, AFD, [&] {
         CounterExpr &funcCounter = assignCounter(AFD->getBody());
 
-        if (isa<ConstructorDecl>(AFD))
+        if (ParentNominalType && isa<ConstructorDecl>(AFD))
           addToCounter(ParentNominalType, funcCounter);
       });
     } else if (auto *TLCD = dyn_cast<TopLevelCodeDecl>(D)) {
@@ -995,16 +992,6 @@ static StringRef getCurrentFileName(ASTNode Root) {
   return {};
 }
 
-static void walkTopLevelNodeForProfiling(ASTNode Root, ASTWalker &Walker) {
-  Root.walk(Walker);
-
-  // Visit extensions when walking through a nominal type.
-  auto *NTD = dyn_cast_or_null<NominalTypeDecl>(Root.dyn_cast<Decl *>());
-  if (NTD)
-    for (ExtensionDecl *ED : NTD->getExtensions())
-      ED->walk(Walker);
-}
-
 void SILProfiler::assignRegionCounters() {
   const auto &SM = M.getASTContext().SourceMgr;
 
@@ -1040,15 +1027,15 @@ void SILProfiler::assignRegionCounters() {
       CurrentFuncName, getEquivalentPGOLinkage(CurrentFuncLinkage),
       CurrentFileName);
 
-  walkTopLevelNodeForProfiling(Root, Mapper);
+  Root.walk(Mapper);
 
   NumRegionCounters = Mapper.NextCounter;
   // TODO: Mapper needs to calculate a function hash as it goes.
   PGOFuncHash = 0x0;
 
   if (EmitCoverageMapping) {
     CoverageMapping Coverage(SM);
-    walkTopLevelNodeForProfiling(Root, Coverage);
+    Root.walk(Coverage);
     CovMap =
         Coverage.emitSourceRegions(M, CurrentFuncName, PGOFuncName, PGOFuncHash,
                                    RegionCounterMap, CurrentFileName);
@@ -1066,7 +1053,7 @@ void SILProfiler::assignRegionCounters() {
     }
     PGOMapping pgoMapper(RegionLoadedCounterMap, LoadedCounts,
                          RegionCondToParentMap);
-    walkTopLevelNodeForProfiling(Root, pgoMapper);
+    Root.walk(pgoMapper);
   }
 }
 

lib/SILGen/SILGen.cpp

@@ -489,13 +489,22 @@ static bool haveProfiledAssociatedFunction(SILDeclRef constant) {
 }
 
 SILProfiler *
-SILGenModule::getOrCreateProfilerForConstructors(NominalTypeDecl *decl) {
-  assert(decl && "No nominal type for constructor lookup");
-
+SILGenModule::getOrCreateProfilerForConstructors(DeclContext *ctx,
+                                                 ConstructorDecl *cd) {
   const auto &Opts = M.getOptions();
   if (!Opts.GenerateProfile && Opts.UseProfile.empty())
     return nullptr;
 
+  // Profile nominal types and extensions separately, as they may live in
+  // distinct files. For extensions, just pass in the constructor, because
+  // there are no stored property initializers to visit.
+  Decl *decl = nullptr;
+  if (ctx->isExtensionContext())
+    decl = cd;
+  else
+    decl = ctx->getAsNominalTypeOrNominalTypeExtensionContext();
+  assert(decl && "No decl available for profiling in this context");
+
   SILProfiler *&profiler = constructorProfilers[decl];
   if (!profiler)
     profiler = SILProfiler::create(M, ForDefinition, decl);
@@ -755,7 +764,7 @@ void SILGenModule::emitConstructor(ConstructorDecl *decl) {
 
   bool ForCoverageMapping = doesASTRequireProfiling(M, decl);
 
-  if (auto *clsDecl = declCtx->getAsClassOrClassExtensionContext()) {
+  if (declCtx->getAsClassOrClassExtensionContext()) {
     // Class constructors have separate entry points for allocation and
     // initialization.
     emitOrDelayFunction(
@@ -772,24 +781,24 @@ void SILGenModule::emitConstructor(ConstructorDecl *decl) {
       SILDeclRef initConstant(decl, SILDeclRef::Kind::Initializer);
       emitOrDelayFunction(
           *this, initConstant,
-          [this, initConstant, decl, clsDecl](SILFunction *initF) {
+          [this, initConstant, decl, declCtx](SILFunction *initF) {
             preEmitFunction(initConstant, decl, initF, decl);
             PrettyStackTraceSILFunction X("silgen constructor initializer",
                                           initF);
-            initF->setProfiler(getOrCreateProfilerForConstructors(clsDecl));
+            initF->setProfiler(
+                getOrCreateProfilerForConstructors(declCtx, decl));
             SILGenFunction(*this, *initF).emitClassConstructorInitializer(decl);
             postEmitFunction(initConstant, initF);
           },
           /*forceEmission=*/ForCoverageMapping);
     }
   } else {
     // Struct and enum constructors do everything in a single function.
-    auto *nomDecl = declCtx->getAsNominalTypeOrNominalTypeExtensionContext();
     emitOrDelayFunction(
-        *this, constant, [this, constant, decl, nomDecl](SILFunction *f) {
+        *this, constant, [this, constant, decl, declCtx](SILFunction *f) {
           preEmitFunction(constant, decl, f, decl);
           PrettyStackTraceSILFunction X("silgen emitConstructor", f);
-          f->setProfiler(getOrCreateProfilerForConstructors(nomDecl));
+          f->setProfiler(getOrCreateProfilerForConstructors(declCtx, decl));
           SILGenFunction(*this, *f).emitValueConstructor(decl);
           postEmitFunction(constant, f);
         });
@@ -997,10 +1006,8 @@ emitStoredPropertyInitialization(PatternBindingDecl *pbd, unsigned i) {
     PrettyStackTraceSILFunction X("silgen emitStoredPropertyInitialization", f);
 
     // Inherit a profiler instance from the constructor.
-    auto *nominalDecl =
-        var->getDeclContext()->getAsNominalTypeOrNominalTypeExtensionContext();
-    if (nominalDecl)
-      f->setProfiler(getOrCreateProfilerForConstructors(nominalDecl));
+    f->setProfiler(
+        getOrCreateProfilerForConstructors(var->getDeclContext(), nullptr));
 
     SILGenFunction(*this, *f).emitGeneratorFunction(constant, init);
     postEmitFunction(constant, f);

lib/SILGen/SILGen.h

@@ -96,9 +96,10 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// The most recent conformance...
   NormalProtocolConformance *lastEmittedConformance = nullptr;
 
-  /// Profiler instances for constructors, grouped by associated nominal type.
+  /// Profiler instances for constructors, grouped by associated decl.
   /// Each profiler is shared by all member initializers for a nominal type.
-  llvm::DenseMap<NominalTypeDecl *, SILProfiler *> constructorProfilers;
+  /// Constructors within extensions are profiled separately.
+  llvm::DenseMap<Decl *, SILProfiler *> constructorProfilers;
 
   SILFunction *emitTopLevelFunction(SILLocation Loc);
 
@@ -440,9 +441,12 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// Emit a property descriptor for the given storage decl if it needs one.
   void tryEmitPropertyDescriptor(AbstractStorageDecl *decl);
 
-  /// Get or create the profiler instance for a nominal type's constructors.
-  SILProfiler *getOrCreateProfilerForConstructors(NominalTypeDecl *decl);
-  
+  /// Get or create the shared profiler instance for a type's constructors.
+  /// This takes care to create separate profilers for extensions, which may
+  /// reside in a different file than the one where the base type is defined.
+  SILProfiler *getOrCreateProfilerForConstructors(DeclContext *ctx,
+                                                  ConstructorDecl *cd);
+
 private:
   /// Emit the deallocator for a class that uses the objc allocator.
   void emitObjCAllocatorDestructor(ClassDecl *cd, DestructorDecl *dd);

test/SILGen/coverage_decls.swift

@@ -1,7 +1,25 @@
-// RUN: %target-swift-frontend -Xllvm -sil-full-demangle -profile-generate -profile-coverage-mapping -emit-sorted-sil -emit-sil -module-name coverage_decls -primary-file %s %S/Inputs/coverage_imports.swift | %FileCheck %s
+// RUN: %target-swift-frontend -Xllvm -sil-full-demangle -profile-generate -profile-coverage-mapping -emit-sorted-sil -emit-sil -module-name coverage_decls %s %S/Inputs/coverage_imports.swift | %FileCheck %s -check-prefix=ALL
+// RUN: %target-swift-frontend -Xllvm -sil-full-demangle -profile-generate -profile-coverage-mapping -emit-sorted-sil -emit-sil -module-name coverage_decls -primary-file %s %S/Inputs/coverage_imports.swift | %FileCheck %s -check-prefix=PRIMARY
 
-// CHECK: sil_coverage_map {{.*}} $S14coverage_decls4mainyyF
-// CHECK-NOT: sil_coverage_map
+// ALL: sil_coverage_map {{.*}} // closure #1 () -> Swift.Int in coverage_decls.Box.x.getter : Swift.Int
+// ALL: sil_coverage_map {{.*}} // coverage_decls.Box.init(y: Swift.Int) -> coverage_decls.Box
+// ALL: sil_coverage_map {{.*}} // coverage_decls.Box.init(z: Swift.String) -> coverage_decls.Box
+// ALL: sil_coverage_map {{.*}} // coverage_decls.main() -> ()
+// ALL: sil_coverage_map {{.*}} // __ntd_Box
+
+// PRIMARY-NOT: sil_coverage_map
+// PRIMARY: sil_coverage_map {{.*}} // coverage_decls.Box.init(y: Swift.Int) -> coverage_decls.Box
+// PRIMARY: sil_coverage_map {{.*}} // coverage_decls.Box.init(z: Swift.String) -> coverage_decls.Box
+// PRIMARY: sil_coverage_map {{.*}} // coverage_decls.main() -> ()
+// PRIMARY-NOT: sil_coverage_map
+
+extension Box {
+  init(y: Int) { self.init() }
+}
+
+extension Box {
+  init(z: String) { self.init() }
+}
 
 func main() {
   var b = Box()