[Profiler] Support lazy variable initializers

Start visiting LazyInitializerExpr for profiling,
such that we emit a profile counter when
initializing the initial value for the first time.

rdar://43393937

Author: hamishknight

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);
 }

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
+}

test/Profiler/pgo_lazy.swift

@@ -0,0 +1,41 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %s -profile-generate -Xfrontend -disable-incremental-llvm-codegen -module-name pgo_lazy -o %t/main
+
+// This unusual use of 'sh' allows the path of the profraw file to be
+// substituted by %target-run.
+// RUN: %target-codesign %t/main
+// RUN: %target-run sh -c 'env LLVM_PROFILE_FILE=$1 $2' -- %t/default.profraw %t/main
+
+// RUN: %llvm-profdata merge %t/default.profraw -o %t/default.profdata
+
+// RUN: %target-swift-frontend %s -Xllvm -sil-full-demangle -profile-use=%t/default.profdata -emit-sorted-sil -emit-sil -module-name pgo_lazy -o - | %FileCheck %s --check-prefix=SIL
+// RUN: %target-swift-frontend %s -Xllvm -sil-full-demangle -profile-use=%t/default.profdata -O -emit-sorted-sil -emit-sil -module-name pgo_lazy -o - | %FileCheck %s --check-prefix=SIL
+// RUN: %target-swift-frontend %s -Xllvm -sil-full-demangle -profile-use=%t/default.profdata -emit-ir -module-name pgo_lazy -o - | %FileCheck %s --check-prefix=IR
+// RUN: %target-swift-frontend %s -Xllvm -sil-full-demangle -profile-use=%t/default.profdata -O -emit-ir -module-name pgo_lazy -o - | %FileCheck %s --check-prefix=IR
+
+// REQUIRES: profile_runtime
+// REQUIRES: executable_test
+// REQUIRES: OS=macosx
+
+var cond = true
+
+public struct S {
+  // SIL-LABEL: sil [lazy_getter] [noinline] @$s8pgo_lazy1SV1xSivg : $@convention(method) (@inout S) -> Int !function_entry_count(35)
+  // SIL: cond_br {{.*}} !true_count(5) !false_count(30)
+  public lazy var x = cond ? 2 : 3
+}
+
+func triggerLazy() -> Int {
+  var s = S()
+  return s.x
+}
+public var total = 0
+for _ in 0 ..< 5 {
+  total += triggerLazy()
+}
+cond = false
+for _ in 0 ..< 30 {
+  total += triggerLazy()
+}
+
+// IR: !{!"branch_weights", i32 6, i32 31}