Merge pull request #60643 from hamishknight/new-coverage-just-dropped

Author: hamishknight

include/swift/AST/ASTWalker.h

@@ -50,6 +50,25 @@ struct ReferenceMetaData {
       : Kind(Kind), AccKind(AccKind), isImplicit(isImplicit) {}
 };
 
+/// Specifies how the initialization expression of a \c lazy variable should be
+/// walked by the ASTWalker.
+enum class LazyInitializerWalking {
+  /// No lazy initialization expressions will be walked.
+  None,
+
+  /// The lazy initialization expression will only be walked as a part of
+  /// the variable's pattern binding decl. This is the default behavior, and is
+  /// consistent with the initializer being syntactically part of the pattern
+  /// binding.
+  InPatternBinding,
+
+  /// The lazy initialization expression will only be walked as part of the
+  /// body of the synthesized accessor for the lazy variable. In such an
+  /// accessor, the expression is denoted by LazyInitializerExpr. This is mainly
+  /// useful for code emission.
+  InAccessor
+};
+
 /// An abstract class used to traverse an AST.
 class ASTWalker {
 public:
@@ -193,15 +212,13 @@ class ASTWalker {
   /// params in AbstractFunctionDecl and NominalTypeDecl.
   virtual bool shouldWalkIntoGenericParams() { return false; }
 
-  /// This method configures whether the walker should walk into the
-  /// initializers of lazy variables.  These initializers are semantically
-  /// different from other initializers in their context and so sometimes
-  /// should not be visited.
-  ///
-  /// Note that visiting the body of the lazy getter will find a
-  /// LazyInitializerExpr with the initializer as its sub-expression.
-  /// However, ASTWalker does not walk into LazyInitializerExprs on its own.
-  virtual bool shouldWalkIntoLazyInitializers() { return true; }
+  /// This method configures how the walker should walk the initializers of
+  /// lazy variables. These initializers are semantically different from other
+  /// initializers in their context and so sometimes should be visited as part
+  /// of the synthesized getter, or should not be visited at all.
+  virtual LazyInitializerWalking getLazyInitializerWalkingBehavior() {
+    return LazyInitializerWalking::InPatternBinding;
+  }
 
   /// This method configures whether the walker should visit the body of a
   /// closure that was checked separately from its enclosing expression.

include/swift/AST/Expr.h

@@ -5314,6 +5314,7 @@ class LazyInitializerExpr : public Expr {
   SourceLoc getLoc() const { return SubExpr->getLoc(); }
 
   Expr *getSubExpr() const { return SubExpr; }
+  void setSubExpr(Expr *subExpr) { SubExpr = subExpr; }
 
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::LazyInitializer;

lib/AST/ASTVerifier.cpp

@@ -2338,11 +2338,11 @@ class Verifier : public ASTWalker {
       verifyCheckedBase(VD);
     }
 
-    bool shouldWalkIntoLazyInitializers() override {
+    LazyInitializerWalking getLazyInitializerWalkingBehavior() override {
       // We don't want to walk into lazy initializers because they should
       // have been reparented to their synthesized getter, which will
       // invalidate various invariants.
-      return false;
+      return LazyInitializerWalking::None;
     }
 
     void verifyChecked(PatternBindingDecl *binding) {

lib/AST/ASTWalker.cpp

@@ -195,18 +195,23 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
         PBD->setPattern(idx, Pat, PBD->getInitContext(idx));
       else
         return true;
-      if (PBD->getInit(idx) &&
-          !isPropertyWrapperBackingProperty &&
-          (!PBD->isInitializerSubsumed(idx) ||
-           Walker.shouldWalkIntoLazyInitializers())) {
+
+      if (!PBD->getInit(idx) || isPropertyWrapperBackingProperty)
+        continue;
+
+      if (PBD->isInitializerSubsumed(idx) &&
+          Walker.getLazyInitializerWalkingBehavior() !=
+              LazyInitializerWalking::InPatternBinding) {
+        break;
+      }
+
 #ifndef NDEBUG
-        PrettyStackTraceDecl debugStack("walking into initializer for", PBD);
+      PrettyStackTraceDecl debugStack("walking into initializer for", PBD);
 #endif
-        if (Expr *E2 = doIt(PBD->getInit(idx)))
-          PBD->setInit(idx, E2);
-        else
-          return true;
-      }
+      if (Expr *E2 = doIt(PBD->getInit(idx)))
+        PBD->setInit(idx, E2);
+      else
+        return true;
     }
     return false;
   }
@@ -1073,7 +1078,17 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   Expr *visitLazyInitializerExpr(LazyInitializerExpr *E) {
-    // Initializer is totally opaque for most visitation purposes.
+    // The initializer is opaque unless we specifically want to visit it as part
+    // of the accessor body.
+    if (Walker.getLazyInitializerWalkingBehavior() !=
+        LazyInitializerWalking::InAccessor) {
+      return E;
+    }
+    auto *sub = doIt(E->getSubExpr());
+    if (!sub)
+      return nullptr;
+
+    E->setSubExpr(sub);
     return E;
   }
 

lib/SIL/IR/SILProfiler.cpp

@@ -196,6 +196,13 @@ static bool skipExpr(Expr *E) {
   return !E->getStartLoc().isValid() || !E->getEndLoc().isValid();
 }
 
+/// Whether the children of an unmapped decl should still be walked.
+static bool shouldWalkUnmappedDecl(const Decl *D) {
+  // We want to walk into the initializer for a pattern binding decl. This
+  // allows us to map LazyInitializerExprs.
+  return isa<PatternBindingDecl>(D);
+}
+
 /// An ASTWalker that maps ASTNodes to profiling counters.
 struct MapRegionCounters : public ASTWalker {
   /// The next counter value to assign.
@@ -207,6 +214,12 @@ struct MapRegionCounters : public ASTWalker {
   MapRegionCounters(llvm::DenseMap<ASTNode, unsigned> &CounterMap)
       : CounterMap(CounterMap) {}
 
+  LazyInitializerWalking getLazyInitializerWalkingBehavior() override {
+    // We want to walk lazy initializers present in the synthesized getter for
+    // a lazy variable.
+    return LazyInitializerWalking::InAccessor;
+  }
+
   void mapRegion(ASTNode N) {
     CounterMap[N] = NextCounter;
 
@@ -225,7 +238,7 @@ struct MapRegionCounters : public ASTWalker {
 
   bool walkToDeclPre(Decl *D) override {
     if (isUnmapped(D))
-      return false;
+      return shouldWalkUnmappedDecl(D);
 
     if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
       return visitFunctionDecl(*this, AFD, [&] { mapRegion(AFD->getBody()); });
@@ -274,14 +287,8 @@ struct MapRegionCounters : public ASTWalker {
       mapRegion(IE->getThenExpr());
     }
 
-    // rdar://42792053
-    // TODO: There's an outstanding issue here with LazyInitializerExpr. A LIE
-    // is copied into the body of a property getter after type-checking (before
-    // coverage). ASTWalker only visits this expression once via the property's
-    // VarDecl, and does not visit it again within the getter. This results in
-    // missing coverage. SILGen treats the init expr as part of the getter, but
-    // its SILProfiler has no information about the init because the LIE isn't
-    // visited here.
+    if (isa<LazyInitializerExpr>(E))
+      mapRegion(E);
 
     return {true, E};
   }
@@ -579,7 +586,7 @@ struct PGOMapping : public ASTWalker {
 
   bool walkToDeclPre(Decl *D) override {
     if (isUnmapped(D))
-      return false;
+      return shouldWalkUnmappedDecl(D);
     if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
       return visitFunctionDecl(*this, AFD, [&] {
         setKnownExecutionCount(AFD->getBody());
@@ -591,6 +598,12 @@ struct PGOMapping : public ASTWalker {
     return true;
   }
 
+  LazyInitializerWalking getLazyInitializerWalkingBehavior() override {
+    // We want to walk lazy initializers present in the synthesized getter for
+    // a lazy variable.
+    return LazyInitializerWalking::InAccessor;
+  }
+
   std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
     unsigned parent = getParentCounter();
     auto parentCount = LoadedCounts.Counts[parent];
@@ -674,6 +687,9 @@ struct PGOMapping : public ASTWalker {
       }
       setExecutionCount(elseExpr, subtract(count, thenCount));
     }
+    if (isa<LazyInitializerExpr>(E))
+      setKnownExecutionCount(E);
+
     return {true, E};
   }
 };
@@ -903,6 +919,12 @@ struct CoverageMapping : public ASTWalker {
 public:
   CoverageMapping(const SourceManager &SM) : SM(SM) {}
 
+  LazyInitializerWalking getLazyInitializerWalkingBehavior() override {
+    // We want to walk lazy initializers present in the synthesized getter for
+    // a lazy variable.
+    return LazyInitializerWalking::InAccessor;
+  }
+
   /// Generate the coverage counter mapping regions from collected
   /// source regions.
   SILCoverageMap *emitSourceRegions(
@@ -930,7 +952,7 @@ struct CoverageMapping : public ASTWalker {
 
   bool walkToDeclPre(Decl *D) override {
     if (isUnmapped(D))
-      return false;
+      return shouldWalkUnmappedDecl(D);
 
     if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
       return visitFunctionDecl(*this, AFD, [&] {
@@ -1124,6 +1146,9 @@ struct CoverageMapping : public ASTWalker {
       }
     }
 
+    if (isa<LazyInitializerExpr>(E))
+      assignCounter(E);
+
     if (hasCounter(E) && !Parent.isNull())
       pushRegion(E);
     return {true, E};

lib/SILGen/SILGenExpr.cpp

@@ -415,7 +415,9 @@ namespace {
       return RValue(SGF, E, SGF.emitAddressOfLValue(E->getSubExpr(),
                                                     std::move(lv)));
     }
-    
+
+    RValue visitLazyInitializerExpr(LazyInitializerExpr *E, SGFContext C);
+
     RValue visitApplyExpr(ApplyExpr *E, SGFContext C);
 
     RValue visitDiscardAssignmentExpr(DiscardAssignmentExpr *E, SGFContext C) {
@@ -708,6 +710,15 @@ tryEmitAsBridgingConversion(SILGenFunction &SGF, Expr *E, bool isExplicit,
   return SGF.emitConvertedRValue(subExpr, conversion, C);
 }
 
+RValue RValueEmitter::visitLazyInitializerExpr(LazyInitializerExpr *E,
+                                               SGFContext C) {
+  // We need to emit a profiler count increment specifically for the lazy
+  // initialization, as we don't want to record an increment for every call to
+  // the getter.
+  SGF.emitProfilerIncrement(E);
+  return visit(E->getSubExpr(), C);
+}
+
 RValue RValueEmitter::visitApplyExpr(ApplyExpr *E, SGFContext C) {
   return SGF.emitApplyExpr(E, C);
 }

lib/SILGen/SILGenFunction.cpp

@@ -1163,9 +1163,8 @@ void SILGenFunction::emitProfilerIncrement(ASTNode N) {
   const auto &RegionCounterMap = SP->getRegionCounterMap();
   auto CounterIt = RegionCounterMap.find(N);
 
-  // TODO: Assert that this cannot happen (rdar://42792053).
-  if (CounterIt == RegionCounterMap.end())
-    return;
+  assert(CounterIt != RegionCounterMap.end() &&
+         "cannot increment non-existent counter");
 
   auto Int32Ty = getLoweredType(BuiltinIntegerType::get(32, C));
   auto Int64Ty = getLoweredType(BuiltinIntegerType::get(64, C));

lib/Sema/TypeCheckCaptures.cpp

@@ -210,11 +210,11 @@ class FindCapturedVars : public ASTWalker {
       checkType(VD->getInterfaceType(), VD->getLoc());
   }
 
-  bool shouldWalkIntoLazyInitializers() override {
+  LazyInitializerWalking getLazyInitializerWalkingBehavior() override {
     // We don't want to walk into lazy initializers because they're not
     // really present at this level.  We'll catch them when processing
     // the getter.
-    return false;
+    return LazyInitializerWalking::None;
   }
 
   std::pair<bool, Expr *> walkToDeclRefExpr(DeclRefExpr *DRE) {

test/Profiler/coverage_class.swift

@@ -52,11 +52,3 @@ struct S2 {
   // CHECK-NEXT: [[@LINE+1]]:17 -> [[@LINE+1]]:27 : 0
   var m1: Int = g1 ? 0 : 1
 }
-
-// Test that the crash from SR-8429 is avoided. Follow-up work is
-// needed to generate the correct coverage mapping here. Coverage for
-// `offset` should be associated with its getter, not the class
-// constructor.
-class C2 {
-  lazy var offset: Int = true ? 30 : 55
-}

test/Profiler/coverage_lazy.swift

@@ -0,0 +1,29 @@
+// RUN: %target-swift-frontend -Xllvm -sil-full-demangle -profile-generate -profile-coverage-mapping -emit-sil -module-name coverage_lazy %s | %FileCheck %s
+// RUN: %target-swift-frontend -profile-generate -profile-coverage-mapping -emit-ir %s
+
+// Test that the crash from SR-8429 is avoided, and that we generate the
+// correct coverage.
+class C {
+  // CHECK-LABEL: sil hidden [lazy_getter] [noinline] @$s13coverage_lazy1CC6offsetSivg : $@convention(method) (@guaranteed C) -> Int
+  // CHECK:       switch_enum {{%[0-9]+}} : $Optional<Int>, case #Optional.some!enumelt: {{bb[0-9]}}, case #Optional.none!enumelt: [[INITBB:bb[0-9]]]
+  // CHECK:       [[INITBB]]
+  // CHECK-NEXT:  string_literal
+  // CHECK-NEXT:  integer_literal $Builtin.Int64, 0
+  // CHECK-NEXT:  integer_literal $Builtin.Int32, 4
+  // CHECK-NEXT:  integer_literal $Builtin.Int32, 2
+  // CHECK-NEXT:  int_instrprof_increment
+  // CHECK:       function_ref @$sSb6randomSbyFZ : $@convention(method) (@thin Bool.Type) -> Bool
+  // CHECK:       cond_br {{%[0-9]+}}, [[TRUEBB:bb[0-9]]], {{bb[0-9]}}
+  // CHECK:       [[TRUEBB]]
+  // CHECK-NEXT:  string_literal
+  // CHECK-NEXT:  integer_literal $Builtin.Int64, 0
+  // CHECK-NEXT:  integer_literal $Builtin.Int32, 4
+  // CHECK-NEXT:  integer_literal $Builtin.Int32, 3
+
+  // CHECK-LABEL: sil_coverage_map {{.*}} // coverage_lazy.C.offset.getter : Swift.Int
+  // CHECK-NEXT: [[@LINE+4]]:38 -> [[@LINE+4]]:40 : 3
+  // CHECK-NEXT: [[@LINE+3]]:43 -> [[@LINE+3]]:45 : (2 - 3)
+  // CHECK-NEXT: [[@LINE+2]]:26 -> [[@LINE+2]]:45 : 2
+  // CHECK-NEXT: }
+  lazy var offset: Int = .random() ? 30 : 55
+}