AST: Store the active DeclContext in AvailabilityScope

include/swift/AST/AvailabilityScope.h

@@ -98,6 +98,7 @@ class AvailabilityScope : public ASTAllocated<AvailabilityScope> {
   /// Represents the AST node that introduced an availability scope.
   class IntroNode {
     Reason IntroReason;
+    const DeclContext *DC;
     union {
       SourceFile *SF;
       Decl *D;
@@ -108,25 +109,26 @@ class AvailabilityScope : public ASTAllocated<AvailabilityScope> {
     };
 
   public:
-    IntroNode(SourceFile *SF) : IntroReason(Reason::Root), SF(SF) {}
-    IntroNode(Decl *D, Reason introReason = Reason::Decl)
-        : IntroReason(introReason), D(D) {
-      (void)getAsDecl();    // check that assertion succeeds
-    }
-    IntroNode(IfStmt *IS, bool IsThen)
+    IntroNode(SourceFile *SF);
+    IntroNode(Decl *D, Reason introReason = Reason::Decl);
+    IntroNode(IfStmt *IS, const DeclContext *DC, bool IsThen)
         : IntroReason(IsThen ? Reason::IfStmtThenBranch
                              : Reason::IfStmtElseBranch),
-          IS(IS) {}
-    IntroNode(PoundAvailableInfo *PAI)
-        : IntroReason(Reason::ConditionFollowingAvailabilityQuery), PAI(PAI) {}
-    IntroNode(GuardStmt *GS, bool IsFallthrough)
+          DC(DC), IS(IS) {}
+    IntroNode(PoundAvailableInfo *PAI, const DeclContext *DC)
+        : IntroReason(Reason::ConditionFollowingAvailabilityQuery), DC(DC),
+          PAI(PAI) {}
+    IntroNode(GuardStmt *GS, const DeclContext *DC, bool IsFallthrough)
         : IntroReason(IsFallthrough ? Reason::GuardStmtFallthrough
                                     : Reason::GuardStmtElseBranch),
-          GS(GS) {}
-    IntroNode(WhileStmt *WS) : IntroReason(Reason::WhileStmtBody), WS(WS) {}
+          DC(DC), GS(GS) {}
+    IntroNode(WhileStmt *WS, const DeclContext *DC)
+        : IntroReason(Reason::WhileStmtBody), DC(DC), WS(WS) {}
 
     Reason getReason() const { return IntroReason; }
 
+    const DeclContext *getDeclContext() const { return DC; }
+
     SourceFile *getAsSourceFile() const {
       assert(IntroReason == Reason::Root);
       return SF;
@@ -200,37 +202,39 @@ class AvailabilityScope : public ASTAllocated<AvailabilityScope> {
 
   /// Create an availability scope for the Then branch of the given IfStmt.
   static AvailabilityScope *createForIfStmtThen(ASTContext &Ctx, IfStmt *S,
+                                                const DeclContext *DC,
                                                 AvailabilityScope *Parent,
                                                 const AvailabilityContext Info);
 
   /// Create an availability scope for the Else branch of the given IfStmt.
   static AvailabilityScope *createForIfStmtElse(ASTContext &Ctx, IfStmt *S,
+                                                const DeclContext *DC,
                                                 AvailabilityScope *Parent,
                                                 const AvailabilityContext Info);
 
   /// Create an availability scope for the true-branch control flow to
   /// further StmtConditionElements following a #available() query in
   /// a StmtCondition.
-  static AvailabilityScope *
-  createForConditionFollowingQuery(ASTContext &Ctx, PoundAvailableInfo *PAI,
-                                   const StmtConditionElement &LastElement,
-                                   AvailabilityScope *Parent,
-                                   const AvailabilityContext Info);
+  static AvailabilityScope *createForConditionFollowingQuery(
+      ASTContext &Ctx, PoundAvailableInfo *PAI,
+      const StmtConditionElement &LastElement, const DeclContext *DC,
+      AvailabilityScope *Parent, const AvailabilityContext Info);
 
   /// Create an availability scope for the fallthrough of a GuardStmt.
   static AvailabilityScope *createForGuardStmtFallthrough(
       ASTContext &Ctx, GuardStmt *RS, BraceStmt *ContainingBraceStmt,
-      AvailabilityScope *Parent, const AvailabilityContext Info);
+      const DeclContext *DC, AvailabilityScope *Parent,
+      const AvailabilityContext Info);
 
   /// Create an availability scope for the else branch of a GuardStmt.
   static AvailabilityScope *
-  createForGuardStmtElse(ASTContext &Ctx, GuardStmt *RS,
+  createForGuardStmtElse(ASTContext &Ctx, GuardStmt *RS, const DeclContext *DC,
                          AvailabilityScope *Parent,
                          const AvailabilityContext Info);
 
   /// Create an availability scope for the body of a WhileStmt.
   static AvailabilityScope *
-  createForWhileStmtBody(ASTContext &Ctx, WhileStmt *WS,
+  createForWhileStmtBody(ASTContext &Ctx, WhileStmt *WS, const DeclContext *DC,
                          AvailabilityScope *Parent,
                          const AvailabilityContext Info);
 

lib/AST/AvailabilityScope.cpp

@@ -30,6 +30,14 @@
 
 using namespace swift;
 
+AvailabilityScope::IntroNode::IntroNode(SourceFile *SF)
+    : IntroReason(Reason::Root), DC(SF), SF(SF) {}
+
+AvailabilityScope::IntroNode::IntroNode(Decl *D, Reason introReason)
+    : IntroReason(introReason), DC(D->getDeclContext()), D(D) {
+  (void)getAsDecl(); // check that assertion succeeds
+}
+
 AvailabilityScope::AvailabilityScope(ASTContext &Ctx, IntroNode Node,
                                      AvailabilityScope *Parent,
                                      SourceRange SrcRange,
@@ -99,67 +107,66 @@ AvailabilityScope *AvailabilityScope::createForDeclImplicit(
                                      Parent, SrcRange, Info);
 }
 
-AvailabilityScope *
-AvailabilityScope::createForIfStmtThen(ASTContext &Ctx, IfStmt *S,
-                                       AvailabilityScope *Parent,
-                                       const AvailabilityContext Info) {
+AvailabilityScope *AvailabilityScope::createForIfStmtThen(
+    ASTContext &Ctx, IfStmt *S, const DeclContext *DC,
+    AvailabilityScope *Parent, const AvailabilityContext Info) {
   assert(S);
   assert(Parent);
-  return new (Ctx) AvailabilityScope(Ctx, IntroNode(S, /*IsThen=*/true), Parent,
-                                     S->getThenStmt()->getSourceRange(), Info);
+  return new (Ctx)
+      AvailabilityScope(Ctx, IntroNode(S, DC, /*IsThen=*/true), Parent,
+                        S->getThenStmt()->getSourceRange(), Info);
 }
 
-AvailabilityScope *
-AvailabilityScope::createForIfStmtElse(ASTContext &Ctx, IfStmt *S,
-                                       AvailabilityScope *Parent,
-                                       const AvailabilityContext Info) {
+AvailabilityScope *AvailabilityScope::createForIfStmtElse(
+    ASTContext &Ctx, IfStmt *S, const DeclContext *DC,
+    AvailabilityScope *Parent, const AvailabilityContext Info) {
   assert(S);
   assert(Parent);
   return new (Ctx)
-      AvailabilityScope(Ctx, IntroNode(S, /*IsThen=*/false), Parent,
+      AvailabilityScope(Ctx, IntroNode(S, DC, /*IsThen=*/false), Parent,
                         S->getElseStmt()->getSourceRange(), Info);
 }
 
 AvailabilityScope *AvailabilityScope::createForConditionFollowingQuery(
     ASTContext &Ctx, PoundAvailableInfo *PAI,
-    const StmtConditionElement &LastElement, AvailabilityScope *Parent,
-    const AvailabilityContext Info) {
+    const StmtConditionElement &LastElement, const DeclContext *DC,
+    AvailabilityScope *Parent, const AvailabilityContext Info) {
   assert(PAI);
   assert(Parent);
   SourceRange Range(PAI->getEndLoc(), LastElement.getEndLoc());
-  return new (Ctx) AvailabilityScope(Ctx, PAI, Parent, Range, Info);
+  return new (Ctx)
+      AvailabilityScope(Ctx, IntroNode(PAI, DC), Parent, Range, Info);
 }
 
 AvailabilityScope *AvailabilityScope::createForGuardStmtFallthrough(
     ASTContext &Ctx, GuardStmt *RS, BraceStmt *ContainingBraceStmt,
-    AvailabilityScope *Parent, const AvailabilityContext Info) {
+    const DeclContext *DC, AvailabilityScope *Parent,
+    const AvailabilityContext Info) {
   assert(RS);
   assert(ContainingBraceStmt);
   assert(Parent);
   SourceRange Range(RS->getEndLoc(), ContainingBraceStmt->getEndLoc());
-  return new (Ctx) AvailabilityScope(Ctx, IntroNode(RS, /*IsFallthrough=*/true),
-                                     Parent, Range, Info);
+  return new (Ctx) AvailabilityScope(
+      Ctx, IntroNode(RS, DC, /*IsFallthrough=*/true), Parent, Range, Info);
 }
 
-AvailabilityScope *
-AvailabilityScope::createForGuardStmtElse(ASTContext &Ctx, GuardStmt *RS,
-                                          AvailabilityScope *Parent,
-                                          const AvailabilityContext Info) {
+AvailabilityScope *AvailabilityScope::createForGuardStmtElse(
+    ASTContext &Ctx, GuardStmt *RS, const DeclContext *DC,
+    AvailabilityScope *Parent, const AvailabilityContext Info) {
   assert(RS);
   assert(Parent);
   return new (Ctx)
-      AvailabilityScope(Ctx, IntroNode(RS, /*IsFallthrough=*/false), Parent,
+      AvailabilityScope(Ctx, IntroNode(RS, DC, /*IsFallthrough=*/false), Parent,
                         RS->getBody()->getSourceRange(), Info);
 }
 
-AvailabilityScope *
-AvailabilityScope::createForWhileStmtBody(ASTContext &Ctx, WhileStmt *S,
-                                          AvailabilityScope *Parent,
-                                          const AvailabilityContext Info) {
+AvailabilityScope *AvailabilityScope::createForWhileStmtBody(
+    ASTContext &Ctx, WhileStmt *S, const DeclContext *DC,
+    AvailabilityScope *Parent, const AvailabilityContext Info) {
   assert(S);
   assert(Parent);
-  return new (Ctx)
-      AvailabilityScope(Ctx, S, Parent, S->getBody()->getSourceRange(), Info);
+  return new (Ctx) AvailabilityScope(Ctx, IntroNode(S, DC), Parent,
+                                     S->getBody()->getSourceRange(), Info);
 }
 
 void AvailabilityScope::addChild(AvailabilityScope *Child, ASTContext &Ctx) {

lib/Sema/TypeCheckAvailability.cpp

@@ -412,10 +412,17 @@ class AvailabilityScopeBuilder : private ASTWalker {
   };
   std::vector<DeclBodyContextInfo> DeclBodyContextStack;
 
+  std::vector<const DeclContext *> DeclContextStack;
+
   AvailabilityScope *getCurrentScope() {
     return ContextStack.back().Scope;
   }
 
+  const DeclContext *getCurrentDeclContext() const {
+    assert(!DeclContextStack.empty());
+    return DeclContextStack.back();
+  }
+
   bool isCurrentScopeContainedByDeploymentTarget() {
     return ContextStack.back().ContainedByDeploymentTarget;
   }
@@ -467,6 +474,7 @@ class AvailabilityScopeBuilder : private ASTWalker {
       : Context(Context) {
     assert(Scope);
     pushContext(Scope, ParentTy());
+    DeclContextStack.push_back(Scope->getIntroductionNode().getDeclContext());
   }
 
   void build(Decl *D) {
@@ -588,27 +596,30 @@ class AvailabilityScopeBuilder : private ASTWalker {
 
     // Implicit decls don't have source locations so they cannot have a scope.
     // However, some implicit nodes contain non-implicit nodes (e.g. defer
-    // blocks) so continue rather than skipping the node entirely.
-    if (D->isImplicit())
-      return Action::Continue();
-
-    if (shouldSkipDecl(D))
-      return Action::SkipNode();
-
-    // The AST of this decl may not be ready to traverse yet if it hasn't been
-    // full typechecked. If that's the case, we leave a placeholder node in the
-    // tree to indicate that the subtree should be expanded lazily when it
-    // needs to be traversed.
-    if (buildLazyContextForDecl(D))
-      return Action::SkipNode();
+    // blocks) so we must continue through them.
+    if (!D->isImplicit()) {
+      if (shouldSkipDecl(D))
+        return Action::SkipNode();
+
+      // The AST of this decl may not be ready to traverse yet if it hasn't been
+      // full typechecked. If that's the case, we leave a placeholder node in
+      // the tree to indicate that the subtree should be expanded lazily when it
+      // needs to be traversed.
+      if (buildLazyContextForDecl(D))
+        return Action::SkipNode();
+
+      // Adds in a scope that covers the entire declaration.
+      if (auto DeclScope = getNewContextForSignatureOfDecl(D)) {
+        pushContext(DeclScope, D);
+      }
 
-    // Adds in a scope that covers the entire declaration.
-    if (auto DeclScope = getNewContextForSignatureOfDecl(D)) {
-      pushContext(DeclScope, D);
+      // Create scopes that cover only the body of the declaration.
+      buildContextsForBodyOfDecl(D);
     }
 
-    // Create scopes that cover only the body of the declaration.
-    buildContextsForBodyOfDecl(D);
+    if (auto *DC = dyn_cast<DeclContext>(D)) {
+      DeclContextStack.push_back(DC);
+    }
 
     // If this decl is the concrete syntax decl for some abstract syntax decl,
     // push it onto the stack so that the abstract syntax decls may be visited.
@@ -624,6 +635,11 @@ class AvailabilityScopeBuilder : private ASTWalker {
       ConcreteDeclStack.pop_back();
     }
 
+    if (auto *DC = dyn_cast<DeclContext>(D)) {
+      assert(DeclContextStack.back() == DC);
+      DeclContextStack.pop_back();
+    }
+
     while (ContextStack.back().ScopeNode.getAsDecl() == D) {
       ContextStack.pop_back();
     }
@@ -1041,7 +1057,8 @@ class AvailabilityScopeBuilder : private ASTWalker {
       auto AvailabilityContext =
           constrainCurrentAvailabilityWithPlatformRange(ThenRange.value());
       auto *ThenScope = AvailabilityScope::createForIfStmtThen(
-          Context, IS, getCurrentScope(), AvailabilityContext);
+          Context, IS, getCurrentDeclContext(), getCurrentScope(),
+          AvailabilityContext);
       AvailabilityScopeBuilder(ThenScope, Context).build(IS->getThenStmt());
     } else {
       build(IS->getThenStmt());
@@ -1064,7 +1081,8 @@ class AvailabilityScopeBuilder : private ASTWalker {
       auto AvailabilityContext =
           constrainCurrentAvailabilityWithPlatformRange(ElseRange.value());
       auto *ElseScope = AvailabilityScope::createForIfStmtElse(
-          Context, IS, getCurrentScope(), AvailabilityContext);
+          Context, IS, getCurrentDeclContext(), getCurrentScope(),
+          AvailabilityContext);
       AvailabilityScopeBuilder(ElseScope, Context).build(ElseStmt);
     } else {
       build(IS->getElseStmt());
@@ -1085,7 +1103,8 @@ class AvailabilityScopeBuilder : private ASTWalker {
       auto AvailabilityContext =
           constrainCurrentAvailabilityWithPlatformRange(BodyRange.value());
       auto *BodyScope = AvailabilityScope::createForWhileStmtBody(
-          Context, WS, getCurrentScope(), AvailabilityContext);
+          Context, WS, getCurrentDeclContext(), getCurrentScope(),
+          AvailabilityContext);
       AvailabilityScopeBuilder(BodyScope, Context).build(WS->getBody());
     } else {
       build(WS->getBody());
@@ -1118,7 +1137,8 @@ class AvailabilityScopeBuilder : private ASTWalker {
         auto AvailabilityContext =
             constrainCurrentAvailabilityWithPlatformRange(ElseRange.value());
         auto *TrueScope = AvailabilityScope::createForGuardStmtElse(
-            Context, GS, getCurrentScope(), AvailabilityContext);
+            Context, GS, getCurrentDeclContext(), getCurrentScope(),
+            AvailabilityContext);
 
         AvailabilityScopeBuilder(TrueScope, Context).build(ElseBody);
       } else {
@@ -1135,7 +1155,8 @@ class AvailabilityScopeBuilder : private ASTWalker {
     auto FallthroughAvailability =
         constrainCurrentAvailabilityWithPlatformRange(FallthroughRange.value());
     auto *FallthroughScope = AvailabilityScope::createForGuardStmtFallthrough(
-        Context, GS, ParentBrace, getCurrentScope(), FallthroughAvailability);
+        Context, GS, ParentBrace, getCurrentDeclContext(), getCurrentScope(),
+        FallthroughAvailability);
 
     pushContext(FallthroughScope, ParentBrace);
   }
@@ -1293,7 +1314,8 @@ class AvailabilityScopeBuilder : private ASTWalker {
       auto ConstrainedAvailability =
           constrainCurrentAvailabilityWithPlatformRange(NewConstraint);
       auto *Scope = AvailabilityScope::createForConditionFollowingQuery(
-          Context, Query, LastElement, CurrentScope, ConstrainedAvailability);
+          Context, Query, LastElement, getCurrentDeclContext(), CurrentScope,
+          ConstrainedAvailability);
 
       pushContext(Scope, ParentTy());
       ++NestedCount;
@@ -1399,6 +1421,11 @@ class AvailabilityScopeBuilder : private ASTWalker {
 
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
     (void)consumeDeclBodyContextIfNecessary(E);
+
+    if (auto CE = dyn_cast<ClosureExpr>(E)) {
+      DeclContextStack.push_back(CE);
+    }
+
     return Action::Continue(E);
   }
 
@@ -1407,6 +1434,11 @@ class AvailabilityScopeBuilder : private ASTWalker {
       ContextStack.pop_back();
     }
 
+    if (auto *CE = dyn_cast<ClosureExpr>(E)) {
+      assert(DeclContextStack.back() == CE);
+      DeclContextStack.pop_back();
+    }
+
     return Action::Continue(E);
   }
 };