Merge pull request #61580 from hamishknight/level-up

Author: hamishknight

include/swift/SIL/SILDeclRef.h

@@ -565,6 +565,11 @@ struct SILDeclRef {
   /// subclassed.
   SubclassScope getSubclassScope() const;
 
+  /// For a SILDeclRef that describes a variable initializer or backing
+  /// initializer, retrieves the expression that will be emitted for that
+  /// initialization. Otherwise, returns \c nullptr.
+  Expr *getInitializationExpr() const;
+
   bool isDynamicallyReplaceable() const;
 
   bool canBeDynamicReplacement() const;

include/swift/SIL/SILFunction.h

@@ -576,9 +576,7 @@ class SILFunction
     Profiler = InheritedProfiler;
   }
 
-  void createProfiler(ASTNode Root, SILDeclRef Ref);
-
-  void discardProfiler() { Profiler = nullptr; }
+  void createProfiler(SILDeclRef Ref);
 
   ProfileCounter getEntryCount() const { return EntryCount; }
 

include/swift/SIL/SILProfiler.h

@@ -83,8 +83,6 @@ class SILProfiler : public SILAllocated<SILProfiler> {
 private:
   SILModule &M;
 
-  ASTNode Root;
-
   SILDeclRef forDecl;
 
   bool EmitCoverageMapping;
@@ -107,13 +105,11 @@ class SILProfiler : public SILAllocated<SILProfiler> {
 
   std::vector<std::tuple<std::string, uint64_t, std::string>> CoverageData;
 
-  SILProfiler(SILModule &M, ASTNode Root, SILDeclRef forDecl,
-              bool EmitCoverageMapping)
-      : M(M), Root(Root), forDecl(forDecl),
-        EmitCoverageMapping(EmitCoverageMapping) {}
+  SILProfiler(SILModule &M, SILDeclRef forDecl, bool EmitCoverageMapping)
+      : M(M), forDecl(forDecl), EmitCoverageMapping(EmitCoverageMapping) {}
 
 public:
-  static SILProfiler *create(SILModule &M, ASTNode N, SILDeclRef Ref);
+  static SILProfiler *create(SILModule &M, SILDeclRef Ref);
 
   /// Check if the function is set up for profiling.
   bool hasRegionCounters() const { return NumRegionCounters != 0; }

lib/SIL/IR/SILDeclRef.cpp

@@ -16,6 +16,7 @@
 #include "swift/AST/AnyFunctionRef.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/ParameterList.h"
+#include "swift/AST/PropertyWrappers.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/ClangImporter/ClangImporter.h"
 #include "swift/ClangImporter/ClangModule.h"
@@ -1497,6 +1498,45 @@ SubclassScope SILDeclRef::getSubclassScope() const {
   llvm_unreachable("Unhandled access level in switch.");
 }
 
+Expr *SILDeclRef::getInitializationExpr() const {
+  switch (kind) {
+  case Kind::StoredPropertyInitializer: {
+    auto *var = cast<VarDecl>(getDecl());
+    auto *pbd = var->getParentPatternBinding();
+    unsigned idx = pbd->getPatternEntryIndexForVarDecl(var);
+    auto *init = pbd->getInit(idx);
+    assert(!pbd->isInitializerSubsumed(idx));
+
+    // If this is the backing storage for a property with an attached wrapper
+    // that was initialized with `=`, use that expression as the initializer.
+    if (auto originalProperty = var->getOriginalWrappedProperty()) {
+      if (originalProperty->isPropertyMemberwiseInitializedWithWrappedType()) {
+        auto wrapperInfo =
+            originalProperty->getPropertyWrapperInitializerInfo();
+        auto *placeholder = wrapperInfo.getWrappedValuePlaceholder();
+        init = placeholder->getOriginalWrappedValue();
+        assert(init);
+      }
+    }
+    return init;
+  }
+  case Kind::PropertyWrapperBackingInitializer: {
+    auto *var = cast<VarDecl>(getDecl());
+    auto wrapperInfo = var->getPropertyWrapperInitializerInfo();
+    assert(wrapperInfo.hasInitFromWrappedValue());
+    return wrapperInfo.getInitFromWrappedValue();
+  }
+  case Kind::PropertyWrapperInitFromProjectedValue: {
+    auto *var = cast<VarDecl>(getDecl());
+    auto wrapperInfo = var->getPropertyWrapperInitializerInfo();
+    assert(wrapperInfo.hasInitFromProjectedValue());
+    return wrapperInfo.getInitFromProjectedValue();
+  }
+  default:
+    return nullptr;
+  }
+}
+
 unsigned SILDeclRef::getParameterListCount() const {
   // Only decls can introduce currying.
   if (!hasDecl())

lib/SIL/IR/SILFunction.cpp

@@ -332,12 +332,11 @@ void SILFunction::deleteSnapshot(int ID) {
   } while ((f = f->snapshots) != nullptr);
 }
 
-void SILFunction::createProfiler(ASTNode Root, SILDeclRef Ref) {
+void SILFunction::createProfiler(SILDeclRef Ref) {
   assert(!Profiler && "Function already has a profiler");
-  assert(Root && "Cannot profile a null ASTNode");
   assert(Ref && "Must have non-null SILDeclRef");
 
-  Profiler = SILProfiler::create(Module, Root, Ref);
+  Profiler = SILProfiler::create(Module, Ref);
   if (!Profiler)
     return;
 

lib/SIL/IR/SILProfiler.cpp

@@ -25,13 +25,75 @@
 
 using namespace swift;
 
-/// Check whether a root AST node should be profiled.
-static bool shouldProfile(ASTNode N, SILDeclRef Constant) {
+/// Unfortunately this is needed as ASTNode can't currently represent a
+/// SourceFile.
+class NodeToProfile final {
+  /// For a direct ASTNode, this stores the node itself. For a main SourceFile,
+  /// it stores the corresponding ModuleDecl.
+  ASTNode Storage;
+
+  explicit NodeToProfile(ASTNode Node) : Storage(Node) {}
+
+public:
+  static NodeToProfile node(ASTNode Node) {
+    assert(!isa_and_nonnull<ModuleDecl>(Node.dyn_cast<Decl *>()));
+    return NodeToProfile(Node);
+  }
+  static NodeToProfile mainSourceFile(SourceFile *SF) {
+    assert(SF->isScriptMode());
+    auto N = NodeToProfile(SF->getParentModule());
+    assert(N.getAsSourceFile() == SF);
+    return N;
+  }
+
+  /// If an ASTNode is being stored, returns it, otherwise \c nullptr.
+  ASTNode getAsNode() const {
+    return isSourceFile() ? nullptr : Storage;
+  }
+
+  /// Whether this is storing a main SourceFile.
+  bool isSourceFile() const {
+    return getAsSourceFile();
+  }
+
+  /// If a main SourceFile is being stored, returns it, otherwise \c nullptr.
+  SourceFile *getAsSourceFile() const {
+    auto *M = dyn_cast_or_null<ModuleDecl>(Storage.dyn_cast<Decl *>());
+    return M ? &M->getMainSourceFile() : nullptr;
+  }
+};
+
+static NodeToProfile getNodeToProfile(SILDeclRef Constant) {
+  // If we have an initialization expression, walk that instead of the variable.
+  if (auto *E = Constant.getInitializationExpr())
+    return NodeToProfile::node(E);
+
+  // Otherwise, we walk the SILDeclRef's node directly.
+  using LocKind = SILDeclRef::LocKind;
+  switch (Constant.getLocKind()) {
+  case LocKind::Decl:
+    return NodeToProfile::node(Constant.getDecl());
+  case LocKind::Closure:
+    return NodeToProfile::node(Constant.getAbstractClosureExpr());
+  case LocKind::File: {
+    auto *SF = cast<SourceFile>(Constant.getFileUnit());
+    return NodeToProfile::mainSourceFile(SF);
+  }
+  }
+  llvm_unreachable("Unhandled case in switch!");
+}
+
+/// Check whether we should profile a given SILDeclRef.
+static bool shouldProfile(SILDeclRef Constant) {
+  auto Root = getNodeToProfile(Constant);
+
   // Do not profile AST nodes with invalid source locations.
-  if (N.getStartLoc().isInvalid() || N.getEndLoc().isInvalid()) {
-    LLVM_DEBUG(llvm::dbgs()
-               << "Skipping ASTNode: invalid start/end locations\n");
-    return false;
+  if (auto N = Root.getAsNode()) {
+    if (N.getStartLoc().isInvalid() || N.getEndLoc().isInvalid()) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "Skipping ASTNode: invalid start/end locations\n");
+      return false;
+    }
   }
 
   // Do not profile AST nodes in unavailable contexts.
@@ -62,51 +124,17 @@ static Stmt *getProfilerStmtForCase(CaseStmt *caseStmt) {
   llvm_unreachable("invalid parent kind");
 }
 
-/// Check that the input AST has at least been type-checked.
-LLVM_ATTRIBUTE_UNUSED
-static bool hasFileBeenTypeChecked(SILDeclRef forDecl) {
-  auto *SF = forDecl.getInnermostDeclContext()->getParentSourceFile();
-  return SF && SF->ASTStage >= SourceFile::TypeChecked;
-}
-
-/// Check whether a mapped AST node is valid for profiling.
-static bool canCreateProfilerForAST(ASTNode N, SILDeclRef forDecl) {
-  assert(hasFileBeenTypeChecked(forDecl) &&
-         "Cannot use this AST for profiling");
-
-  if (auto *D = N.dyn_cast<Decl *>()) {
-    if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D))
-      return AFD->hasBody();
-
-    if (isa<TopLevelCodeDecl>(D))
-      return true;
-  } else if (N.get<Expr *>()) {
-    if (auto *closure = forDecl.getAbstractClosureExpr())
-      return closure->hasBody();
-    if (forDecl.isStoredPropertyInitializer() ||
-        forDecl.isPropertyWrapperBackingInitializer())
-      return true;
-  }
-  return false;
-}
-
-SILProfiler *SILProfiler::create(SILModule &M, ASTNode N, SILDeclRef Ref) {
+SILProfiler *SILProfiler::create(SILModule &M, SILDeclRef Ref) {
   // If profiling isn't enabled, don't profile anything.
   const auto &Opts = M.getOptions();
   if (!Opts.GenerateProfile && Opts.UseProfile.empty())
     return nullptr;
 
-  if (!shouldProfile(N, Ref))
+  if (!shouldProfile(Ref))
     return nullptr;
 
-  if (!canCreateProfilerForAST(N, Ref)) {
-    N.dump(llvm::errs());
-    llvm_unreachable("Invalid AST node for profiling");
-  }
-
   auto *Buf = M.allocate<SILProfiler>(1);
-  auto *SP =
-      ::new (Buf) SILProfiler(M, N, Ref, Opts.EmitProfileCoverageMapping);
+  auto *SP = ::new (Buf) SILProfiler(M, Ref, Opts.EmitProfileCoverageMapping);
   SP->assignRegionCounters();
   return SP;
 }
@@ -163,6 +191,7 @@ namespace {
 template <typename F>
 ASTWalker::PreWalkAction
 visitFunctionDecl(ASTWalker &Walker, AbstractFunctionDecl *AFD, F Func) {
+  assert(AFD->hasBody());
   if (Walker.Parent.isNull()) {
     Func();
     return ASTWalker::Action::Continue();
@@ -177,7 +206,9 @@ shouldWalkIntoExpr(Expr *E, ASTWalker::ParentTy Parent, SILDeclRef Constant) {
 
   // Profiling for closures should be handled separately. Do not visit
   // closure expressions twice.
-  if (isa<AbstractClosureExpr>(E)) {
+  if (auto *CE = dyn_cast<AbstractClosureExpr>(E)) {
+    assert(CE->hasBody());
+
     // A non-null parent means we have a closure child, which we will visit
     // separately. Even if the parent is null, don't walk into a closure if the
     // SILDeclRef is not for a closure, as it could be for a property
@@ -1307,29 +1338,33 @@ static StringRef getCurrentFileName(SILDeclRef forDecl) {
   return {};
 }
 
+static void walkNode(NodeToProfile Node, ASTWalker &Walker) {
+  if (auto N = Node.getAsNode()) {
+    N.walk(Walker);
+  } else {
+    // We want to walk the SourceFile for a top-level entry point. We will only
+    // assign regions to TopLevelCodeDecls.
+    Node.getAsSourceFile()->walk(Walker);
+  }
+}
+
 void SILProfiler::assignRegionCounters() {
   const auto &SM = M.getASTContext().SourceMgr;
 
   CurrentFileName = getCurrentFileName(forDecl);
 
   MapRegionCounters Mapper(forDecl, RegionCounterMap);
 
-  std::string CurrentFuncName;
-  FormalLinkage CurrentFuncLinkage;
-  if (auto *D = Root.dyn_cast<Decl *>()) {
-    if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
-      CurrentFuncName = forDecl.mangle();
-      CurrentFuncLinkage = getDeclLinkage(AFD);
-    } else {
-      auto *TLCD = cast<TopLevelCodeDecl>(D);
-      llvm::raw_string_ostream OS{CurrentFuncName};
-      OS << "__tlcd_";
-      TLCD->getStartLoc().printLineAndColumn(OS, SM);
-      CurrentFuncLinkage = FormalLinkage::HiddenUnique;
+  auto Root = getNodeToProfile(forDecl);
+
+  auto CurrentFuncName = forDecl.mangle();
+  auto CurrentFuncLinkage = FormalLinkage::HiddenUnique;
+
+  if (auto N = Root.getAsNode()) {
+    if (auto *D = N.dyn_cast<Decl *>()) {
+      if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D))
+        CurrentFuncLinkage = getDeclLinkage(AFD);
     }
-  } else {
-    CurrentFuncName = forDecl.mangle();
-    CurrentFuncLinkage = FormalLinkage::HiddenUnique;
   }
 
   PGOFuncName = llvm::getPGOFuncName(
@@ -1341,15 +1376,15 @@ void SILProfiler::assignRegionCounters() {
 
   LLVM_DEBUG(llvm::dbgs() << "Assigning counters to: " << CurrentFuncName
                           << "\n");
-  Root.walk(Mapper);
+  walkNode(Root, Mapper);
 
   NumRegionCounters = Mapper.NextCounter;
   // TODO: Mapper needs to calculate a function hash as it goes.
   PGOFuncHash = 0x0;
 
   if (EmitCoverageMapping) {
     CoverageMapping Coverage(SM, forDecl);
-    Root.walk(Coverage);
+    walkNode(Root, Coverage);
     CovMap =
         Coverage.emitSourceRegions(M, CurrentFuncName, PGOFuncName, PGOFuncHash,
                                    RegionCounterMap, CurrentFileName);
@@ -1367,7 +1402,7 @@ void SILProfiler::assignRegionCounters() {
     }
     PGOMapping pgoMapper(forDecl, RegionCounterMap, LoadedCounts.get(),
                          RegionLoadedCounterMap, RegionCondToParentMap);
-    Root.walk(pgoMapper);
+    walkNode(Root, pgoMapper);
   }
 }