[AST] Allow configurable lazy initializer walking

Allow ASTWalker subclasses to specify whether
they want to visit lazy variable initializers as
part of the pattern binding, getter body, or not
at all.

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/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) {