[Sema] Remove RecontextualizeClosures

Merge with ContextualizeClosuresAndMacros, and
rename to ContextualizationWalker given that it
re-contextualizes a whole bunch of AST nodes now.
This ensures we correctly handle cases such as
decls in if/switch expressions within autoclosures.

Author: hamishknight

lib/Sema/TypeCheckStmt.cpp

@@ -61,61 +61,57 @@ using namespace swift;
 #define DEBUG_TYPE "TypeCheckStmt"
 
 namespace {
-  /// After forming autoclosures, we must re-parent any closure expressions
-  /// nested inside the autoclosure, because the autoclosure introduces a new
-  /// DeclContext.
-  class ContextualizeClosuresAndMacros : public ASTWalker {
+  /// After forming autoclosures and lazy initializer getters, we must update
+  /// the DeclContexts for any AST nodes that store the DeclContext they're
+  /// within. This includes e.g closures and decls, as well as some other
+  /// expressions, statements, and patterns.
+  class ContextualizationWalker : public ASTWalker {
     DeclContext *ParentDC;
+    bool ForLazyInit;
+
+    ContextualizationWalker(DeclContext *parent, bool lazyInit)
+      : ParentDC(parent), ForLazyInit(lazyInit) {}
+
   public:
-    ContextualizeClosuresAndMacros(DeclContext *parent) : ParentDC(parent) {}
+    static void contextualize(ASTNode node, DeclContext *DC) {
+      node.walk(ContextualizationWalker(DC, /*forLazyInit*/ false));
+    }
+
+    static void contextualizeLazyInit(Expr *E, DeclContext *DC) {
+      E->walk(ContextualizationWalker(DC, /*forLazyInit*/ true));
+    }
 
     MacroWalking getMacroWalkingBehavior() const override {
       return MacroWalking::ArgumentsAndExpansion;
     }
 
+    LazyInitializerWalking getLazyInitializerWalkingBehavior() override {
+      // Don't walk lazy initializers, we walk the getter body specially when
+      // synthesizing.
+      return LazyInitializerWalking::None;
+    }
+
     PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-      if (auto CE = dyn_cast<AutoClosureExpr>(E)) {
+      if (auto *CE = dyn_cast<AbstractClosureExpr>(E)) {
         CE->setParent(ParentDC);
-
-        // Recurse into the autoclosure body with the new ParentDC.
-        auto oldParentDC = ParentDC;
-        ParentDC = CE;
-        CE->getBody()->walk(*this);
-        ParentDC = oldParentDC;
-
-        TypeChecker::computeCaptures(CE);
-        return Action::SkipNode(E);
-      } 
-
-      if (auto CapE = dyn_cast<CaptureListExpr>(E)) {
-        // Capture lists need to be reparented to enclosing autoclosures
-        // and/or initializers of property wrapper backing properties
-        // (because they subsume initializers associated with wrapped
-        // properties).
-        if (isa<AutoClosureExpr>(ParentDC) ||
-            isPropertyWrapperBackingPropertyInitContext(ParentDC)) {
-          for (auto &Cap : CapE->getCaptureList()) {
-            Cap.PBD->setDeclContext(ParentDC);
-            Cap.getVar()->setDeclContext(ParentDC);
-          }
+        // If we're contextualizing a lazy init, we don't need to walk into any
+        // child DeclContexts, they'll already be correctly parented. We don't
+        // need to compute captures either, that'll have been done already.
+        if (!ForLazyInit) {
+          contextualize(CE->getBody(), CE);
+          TypeChecker::computeCaptures(CE);
         }
-      }
-
-      if (auto CE = dyn_cast<ClosureExpr>(E)) {
-        CE->setParent(ParentDC);
-        CE->getBody()->walk(ContextualizeClosuresAndMacros(CE));
-
-        TypeChecker::computeCaptures(CE);
         return Action::SkipNode(E);
       }
 
       // Caller-side default arguments need their @autoclosures checked.
-      if (auto *DAE = dyn_cast<DefaultArgumentExpr>(E))
+      if (auto *DAE = dyn_cast<DefaultArgumentExpr>(E)) {
         if (DAE->isCallerSide() &&
             (DAE->getParamDecl()->isAutoClosure() ||
              (DAE->getParamDecl()->getDefaultArgumentKind() ==
               DefaultArgumentKind::ExpressionMacro)))
           DAE->getCallerSideDefaultExpr()->walk(*this);
+      }
 
       // Macro expansion expressions require a DeclContext as well.
       if (auto macroExpansion = dyn_cast<MacroExpansionExpr>(E)) {
@@ -125,26 +121,59 @@ namespace {
       return Action::Continue(E);
     }
 
-    /// We don't want to recurse into most local declarations.
-    PreWalkAction walkToDeclPre(Decl *D) override {
-      // But we do want to walk into the initializers of local
-      // variables.
-      return Action::VisitNodeIf(isa<PatternBindingDecl>(D));
+    PreWalkResult<Pattern *> walkToPatternPre(Pattern *P) override {
+      // A couple of patterns store DeclContexts.
+      if (auto *EP = dyn_cast<ExprPattern>(P))
+        EP->setDeclContext(ParentDC);
+      if (auto *EP = dyn_cast<EnumElementPattern>(P))
+        EP->setDeclContext(ParentDC);
+
+      return Action::Continue(P);
     }
 
-  private:
-    static bool isPropertyWrapperBackingPropertyInitContext(DeclContext *DC) {
-      auto *init = dyn_cast<PatternBindingInitializer>(DC);
-      if (!init)
-        return false;
+    PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
+      // The ASTWalker doesn't walk the case body variables, contextualize them
+      // ourselves.
+      if (auto *CS = dyn_cast<CaseStmt>(S)) {
+        for (auto *CaseVar : CS->getCaseBodyVariablesOrEmptyArray())
+          CaseVar->setDeclContext(ParentDC);
+      }
+      // A few statements store DeclContexts, update them.
+      if (auto *BS = dyn_cast<BreakStmt>(S))
+        BS->setDeclContext(ParentDC);
+      if (auto *CS = dyn_cast<ContinueStmt>(S))
+        CS->setDeclContext(ParentDC);
+      if (auto *FS = dyn_cast<FallthroughStmt>(S))
+        FS->setDeclContext(ParentDC);
+
+      return Action::Continue(S);
+    }
 
-      if (auto *PB = init->getBinding()) {
-        auto *var = PB->getSingleVar();
-        return var && var->getOriginalWrappedProperty(
-                          PropertyWrapperSynthesizedPropertyKind::Backing);
+    PreWalkAction walkToDeclPre(Decl *D) override {
+      // We may encounter some decls parented outside of a local context, e.g
+      // VarDecls in TopLevelCodeDecls are parented to the file. In such cases,
+      // assume the DeclContext they already have is correct, autoclosures
+      // and lazy var inits cannot be defined in such contexts anyway.
+      if (!D->getDeclContext()->isLocalContext())
+        return Action::SkipNode();
+
+      D->setDeclContext(ParentDC);
+
+      // Auxiliary decls need to have their contexts adjusted too.
+      if (auto *VD = dyn_cast<VarDecl>(D)) {
+        VD->visitAuxiliaryDecls([&](VarDecl *D) {
+          D->setDeclContext(ParentDC);
+        });
       }
 
-      return false;
+      // We don't currently support peer macro declarations in local contexts,
+      // however we don't reject them either; so just to be safe, adjust their
+      // context too.
+      D->visitAuxiliaryDecls([&](Decl *D) {
+        D->setDeclContext(ParentDC);
+      });
+
+      return Action::VisitNodeIf(isa<PatternBindingDecl>(D));
     }
   };
 
@@ -197,20 +226,21 @@ namespace {
 } // end anonymous namespace
 
 void TypeChecker::contextualizeInitializer(Initializer *DC, Expr *E) {
-  ContextualizeClosuresAndMacros CC(DC);
-  E->walk(CC);
+  ContextualizationWalker::contextualize(E, DC);
 }
 
 void TypeChecker::contextualizeCallSideDefaultArgument(DeclContext *DC,
                                                        Expr *E) {
-  ContextualizeClosuresAndMacros CC(DC);
-  E->walk(CC);
+  ContextualizationWalker::contextualize(E, DC);
 }
 
 void TypeChecker::contextualizeTopLevelCode(TopLevelCodeDecl *TLCD) {
-  ContextualizeClosuresAndMacros CC(TLCD);
   if (auto *body = TLCD->getBody())
-    body->walk(CC);
+    ContextualizationWalker::contextualize(body, TLCD);
+}
+
+void TypeChecker::contextualizeLazyInit(Expr *E, DeclContext *DC) {
+  ContextualizationWalker::contextualizeLazyInit(E, DC);
 }
 
 namespace {
@@ -1033,7 +1063,7 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
   /// Type-check an entire function body.
   bool typeCheckBody(BraceStmt *&S) {
     bool HadError = typeCheckStmt(S);
-    S->walk(ContextualizeClosuresAndMacros(DC));
+    ContextualizationWalker::contextualize(S, DC);
     return HadError;
   }
 
@@ -2915,7 +2945,7 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &eval,
         body = *optBody;
         alreadyTypeChecked = true;
 
-        body->walk(ContextualizeClosuresAndMacros(AFD));
+        ContextualizationWalker::contextualize(body, AFD);
       }
     }
   }

lib/Sema/TypeCheckStorage.cpp

@@ -1660,75 +1660,6 @@ synthesizeRead2CoroutineGetterBody(AccessorDecl *getter, ASTContext &ctx) {
   return synthesizeTrivialGetterBody(getter, TargetImpl::Implementation, ctx);
 }
 
-namespace {
-  /// This ASTWalker explores an expression tree looking for expressions (which
-  /// are DeclContext's) and changes their parent DeclContext to NewDC.
-  /// TODO: We ought to consider merging this with
-  /// ContextualizeClosuresAndMacros, or better yet removing it in favor of
-  /// avoiding the recontextualization for lazy vars.
-  class RecontextualizeClosures : public ASTWalker {
-    DeclContext *NewDC;
-  public:
-    RecontextualizeClosures(DeclContext *NewDC) : NewDC(NewDC) {}
-
-    MacroWalking getMacroWalkingBehavior() const override {
-      return MacroWalking::ArgumentsAndExpansion;
-    }
-
-    PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-      // If we find a closure, update its declcontext and do *not* walk into it.
-      if (auto CE = dyn_cast<AbstractClosureExpr>(E)) {
-        CE->setParent(NewDC);
-        return Action::SkipNode(E);
-      }
-
-      return Action::Continue(E);
-    }
-
-    PreWalkResult<Pattern *> walkToPatternPre(Pattern *P) override {
-      if (auto *EP = dyn_cast<ExprPattern>(P))
-        EP->setDeclContext(NewDC);
-      if (auto *EP = dyn_cast<EnumElementPattern>(P))
-        EP->setDeclContext(NewDC);
-
-      return Action::Continue(P);
-    }
-
-    PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
-      // The ASTWalker doesn't walk the case body variables, contextualize them
-      // ourselves.
-      if (auto *CS = dyn_cast<CaseStmt>(S)) {
-        for (auto *CaseVar : CS->getCaseBodyVariablesOrEmptyArray())
-          CaseVar->setDeclContext(NewDC);
-      }
-      // A few statements store DeclContexts, update them.
-      if (auto *BS = dyn_cast<BreakStmt>(S))
-        BS->setDeclContext(NewDC);
-      if (auto *CS = dyn_cast<ContinueStmt>(S))
-        CS->setDeclContext(NewDC);
-      if (auto *FS = dyn_cast<FallthroughStmt>(S))
-        FS->setDeclContext(NewDC);
-
-      return Action::Continue(S);
-    }
-
-    PreWalkAction walkToDeclPre(Decl *D) override {
-      D->setDeclContext(NewDC);
-
-      // Auxiliary decls need to have their contexts adjusted too.
-      if (auto *VD = dyn_cast<VarDecl>(D)) {
-        VD->visitAuxiliaryDecls([&](VarDecl *D) {
-          D->setDeclContext(NewDC);
-        });
-      }
-
-      // Skip walking the children of any Decls that are also DeclContexts,
-      // they will already have the right parent.
-      return Action::SkipNodeIf(isa<DeclContext>(D));
-    }
-  };
-} // end anonymous namespace
-
 /// Synthesize the getter for a lazy property with the specified storage
 /// vardecl.
 static std::pair<BraceStmt *, bool>
@@ -1800,12 +1731,13 @@ synthesizeLazyGetterBody(AccessorDecl *Get, VarDecl *VD, VarDecl *Storage,
     InitValue = new (Ctx) ErrorExpr(SourceRange(), Tmp2VD->getTypeInContext());
   }
 
-  // Recontextualize any closure declcontexts nested in the initializer to
-  // realize that they are in the getter function.
+  // We may have stolen the self decl from a PatternBindingInitializer,
+  // re-parent it here.
   if (Get->hasImplicitSelfDecl())
     Get->getImplicitSelfDecl()->setDeclContext(Get);
 
-  InitValue->walk(RecontextualizeClosures(Get));
+  // Also Recontextualize any nodes nested in the initializer.
+  TypeChecker::contextualizeLazyInit(InitValue, Get);
 
   // Wrap the initializer in a LazyInitializerExpr to avoid walking it twice.
   auto initType = InitValue->getType();

lib/Sema/TypeChecker.h

@@ -779,6 +779,9 @@ void contextualizeInitializer(Initializer *DC, Expr *init);
 void contextualizeCallSideDefaultArgument(DeclContext *DC, Expr *init);
 void contextualizeTopLevelCode(TopLevelCodeDecl *TLCD);
 
+/// Contextualize a lazy initializer expression within its implicit getter.
+void contextualizeLazyInit(Expr *E, DeclContext *DC);
+
 /// Retrieve the default type for the given protocol.
 ///
 /// Some protocols, particularly those that correspond to literals, have

test/SILGen/if_expr.swift

@@ -659,6 +659,21 @@ struct LazyProp {
   }
 }
 
+// https://github.com/swiftlang/swift/issues/75294
+func testAsyncLet(_ x: Int?) async {
+  async let _ = if let i = x { i } else { 0 }
+  // CHECK-LABEL: sil private [ossa] @$s7if_expr12testAsyncLetyySiSgYaFSiyYaYbcfu_ : $@convention(thin) @Sendable @async @substituted <τ_0_0> (Optional<Int>) -> (@out τ_0_0, @error any Error) for <Int>
+
+  async let _ = if let i = x {
+    // CHECK-LABEL: sil private [ossa] @$s7if_expr12testAsyncLetyySiSgYaFSiyYaYbcfu0_ : $@convention(thin) @Sendable @async @substituted <τ_0_0> (Optional<Int>) -> (@out τ_0_0, @error any Error) for <Int>
+    lazy var y = i
+    // CHECK-LABEL: sil private [lazy_getter] [noinline] [ossa] @$s7if_expr12testAsyncLetyySiSgYaFSiyYaYbcfu0_1yL_Sivg : $@convention(thin) (@guaranteed { var Optional<Int> }, Int) -> Int
+    then y
+  } else {
+    0
+  }
+}
+
 func testNestedFallthrough1() throws -> Int {
   let x = if .random() {
     switch Bool.random() {

test/SILGen/switch_expr.swift

@@ -302,6 +302,12 @@ struct TestPropertyInit {
   }
 }
 
+// https://github.com/swiftlang/swift/issues/75294
+func testAsyncLet(_ x: Int?) async {
+  async let _ = switch x { case let i?: i default: 0 }
+  // CHECK-LABEL: sil private [ossa] @$s11switch_expr12testAsyncLetyySiSgYaFSiyYaYbcfu_ : $@convention(thin) @Sendable @async @substituted <τ_0_0> (Optional<Int>) -> (@out τ_0_0, @error any Error) for <Int>
+}
+
 func testNested(_ e: E) throws -> Int {
   switch e {
   case .a: