ASTScope: Simplify representation of closures

Let's use a ClosureParametersScope for all closures, even those
without an 'in' keyword. This eliminates the need for the
ClosureBodyScope and WholeClosureScope.

Also, let's move the lookup of capture list bindings from
CaptureParametersScope to CaptureListScope. This eliminates the
need for CaptureParametersScope to store a reference to the
capture list, which allows us to remove the AbstractClosureScope
base class entirely.

Author: slavapestov

include/swift/AST/ASTScope.h

@@ -1393,9 +1393,7 @@ class ConditionalClausePatternUseScope final : public ASTScopeImpl {
 };
 
 
-/// Capture lists may contain initializer expressions
-/// No local bindings here (other than closures in initializers);
-/// rather include these in the params or body local bindings
+/// Capture lists introduce local bindings.
 class CaptureListScope final : public ASTScopeImpl {
 public:
   CaptureListExpr *const expr;
@@ -1416,63 +1414,17 @@ class CaptureListScope final : public ASTScopeImpl {
   NullablePtr<Expr> getExprIfAny() const override { return expr; }
   Expr *getExpr() const { return expr; }
   NullablePtr<const void> getReferrent() const override;
+  bool lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
+                             DeclConsumer) const override;
 };
 
-// In order for compatibility with existing lookup, closures are represented
-// by multiple scopes: An overall scope (including the part before the "in"
-// and a body scope, including the part after the "in"
-class AbstractClosureScope : public ASTScopeImpl {
+/// For a closure with named parameters, this scope does the local bindings.
+class ClosureParametersScope final : public ASTScopeImpl {
 public:
-  NullablePtr<CaptureListExpr> captureList;
   ClosureExpr *const closureExpr;
 
-  AbstractClosureScope(ClosureExpr *closureExpr,
-                       NullablePtr<CaptureListExpr> captureList)
-      : captureList(captureList), closureExpr(closureExpr) {}
-  virtual ~AbstractClosureScope() {}
-
-  NullablePtr<ClosureExpr> getClosureIfClosureScope() const override;
-  NullablePtr<DeclContext> getDeclContext() const override {
-    return closureExpr;
-  }
-  NullablePtr<const void> addressForPrinting() const override {
-    return closureExpr;
-  }
-};
-
-class WholeClosureScope final : public AbstractClosureScope {
-  const BraceStmt *bodyWhenLastExpanded;
-
-public:
-  WholeClosureScope(ClosureExpr *closureExpr,
-                    NullablePtr<CaptureListExpr> captureList)
-      : AbstractClosureScope(closureExpr, captureList) {}
-  virtual ~WholeClosureScope() {}
-
-protected:
-  ASTScopeImpl *expandSpecifically(ScopeCreator &scopeCreator) override;
-  void beCurrent() override;
-  bool isCurrentIfWasExpanded() const override;
-
-private:
-  void expandAScopeThatDoesNotCreateANewInsertionPoint(ScopeCreator &);
-
-public:
-  std::string getClassName() const override;
-  SourceRange
-  getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
-  NullablePtr<Expr> getExprIfAny() const override { return closureExpr; }
-  Expr *getExpr() const { return closureExpr; }
-  NullablePtr<const void> getReferrent() const override;
-};
-
-/// For a closure with named parameters, this scope does the local bindings.
-/// Absent if no "in".
-class ClosureParametersScope final : public AbstractClosureScope {
-public:
-  ClosureParametersScope(ClosureExpr *closureExpr,
-                         NullablePtr<CaptureListExpr> captureList)
-      : AbstractClosureScope(closureExpr, captureList) {}
+  ClosureParametersScope(ClosureExpr *closureExpr)
+      : closureExpr(closureExpr) {}
   virtual ~ClosureParametersScope() {}
 
   std::string getClassName() const override;
@@ -1486,35 +1438,25 @@ class ClosureParametersScope final : public AbstractClosureScope {
   ///  a capture of self.
   NullablePtr<DeclContext> capturedSelfDC() const override;
 
-protected:
-  ASTScopeImpl *expandSpecifically(ScopeCreator &) override;
-  bool lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
-                             DeclConsumer) const override;
-  Optional<bool> resolveIsCascadingUseForThisScope(
-      Optional<bool> isCascadingUse) const override;
-};
-
-// The body encompasses the code in the closure; the part after the "in" if
-// there is an "in"
-class ClosureBodyScope final : public AbstractClosureScope {
-public:
-  ClosureBodyScope(ClosureExpr *closureExpr,
-                   NullablePtr<CaptureListExpr> captureList)
-      : AbstractClosureScope(closureExpr, captureList) {}
-  virtual ~ClosureBodyScope() {}
+  NullablePtr<ClosureExpr> getClosureIfClosureScope() const override {
+    return closureExpr;
+  }
+  NullablePtr<DeclContext> getDeclContext() const override {
+    return closureExpr;
+  }
+  NullablePtr<Expr> getExprIfAny() const override { return closureExpr; }
+  Expr *getExpr() const { return closureExpr; }
+  NullablePtr<const void> getReferrent() const override;
 
 protected:
   ASTScopeImpl *expandSpecifically(ScopeCreator &scopeCreator) override;
 
 private:
   void expandAScopeThatDoesNotCreateANewInsertionPoint(ScopeCreator &);
 
-public:
-  std::string getClassName() const override;
-  SourceRange
-  getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
-
 protected:
+  bool lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
+                             DeclConsumer) const override;
   Optional<bool> resolveIsCascadingUseForThisScope(
       Optional<bool> isCascadingUse) const override;
 };

lib/AST/ASTScope.cpp

@@ -96,10 +96,6 @@ NullablePtr<ClosureExpr> BraceStmtScope::parentClosureIfAny() const {
 NullablePtr<ClosureExpr> ASTScopeImpl::getClosureIfClosureScope() const {
   return nullptr;
 }
-NullablePtr<ClosureExpr>
-AbstractClosureScope::getClosureIfClosureScope() const {
-  return closureExpr;
-}
 
 // Conservative, because using precise info would be circular
 SourceRange
@@ -231,9 +227,7 @@ DEFINE_GET_CLASS_NAME(PatternEntryInitializerScope)
 DEFINE_GET_CLASS_NAME(ConditionalClauseScope)
 DEFINE_GET_CLASS_NAME(ConditionalClausePatternUseScope)
 DEFINE_GET_CLASS_NAME(CaptureListScope)
-DEFINE_GET_CLASS_NAME(WholeClosureScope)
 DEFINE_GET_CLASS_NAME(ClosureParametersScope)
-DEFINE_GET_CLASS_NAME(ClosureBodyScope)
 DEFINE_GET_CLASS_NAME(TopLevelCodeScope)
 DEFINE_GET_CLASS_NAME(SpecializeAttributeScope)
 DEFINE_GET_CLASS_NAME(DifferentiableAttributeScope)

lib/AST/ASTScopeCreation.cpp

@@ -383,21 +383,51 @@ class ScopeCreator final {
 
   void addExprToScopeTree(Expr *expr, ASTScopeImpl *parent) {
     // Use the ASTWalker to find buried captures and closures
-    forEachClosureIn(expr, [&](NullablePtr<CaptureListExpr> captureList,
-                               ClosureExpr *closureExpr) {
-      ifUniqueConstructExpandAndInsert<WholeClosureScope>(parent, closureExpr,
-                                                          captureList);
-    });
+    ASTScopeAssert(expr,
+                 "If looking for closures, must have an expression to search.");
+
+    /// AST walker that finds top-level closures in an expression.
+    class ClosureFinder : public ASTWalker {
+      ScopeCreator &scopeCreator;
+      ASTScopeImpl *parent;
+
+    public:
+      ClosureFinder(ScopeCreator &scopeCreator, ASTScopeImpl *parent)
+          : scopeCreator(scopeCreator), parent(parent) {}
+
+      std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+        if (auto *closure = dyn_cast<ClosureExpr>(E)) {
+          scopeCreator
+              .ifUniqueConstructExpandAndInsert<ClosureParametersScope>(
+                  parent, closure);
+          return {false, E};
+        }
+        if (auto *capture = dyn_cast<CaptureListExpr>(E)) {
+          scopeCreator
+              .ifUniqueConstructExpandAndInsert<CaptureListScope>(
+                  parent, capture);
+          return {false, E};
+        }
+        return {true, E};
+      }
+      std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
+        if (isa<BraceStmt>(S)) { // closures hidden in here
+          return {true, S};
+        }
+        return {false, S};
+      }
+      std::pair<bool, Pattern *> walkToPatternPre(Pattern *P) override {
+        return {false, P};
+      }
+      bool walkToDeclPre(Decl *D) override { return false; }
+      bool walkToTypeReprPre(TypeRepr *T) override { return false; }
+      bool walkToParameterListPre(ParameterList *PL) override { return false; }
+    };
+
+    expr->walk(ClosureFinder(*this, parent));
   }
 
 private:
-  /// Find all of the (non-nested) closures (and associated capture lists)
-  /// referenced within this expression.
-  void forEachClosureIn(
-      Expr *expr,
-      function_ref<void(NullablePtr<CaptureListExpr>, ClosureExpr *)>
-          foundClosure);
-
   // A safe way to discover this, without creating a circular request.
   // Cannot call getAttachedPropertyWrappers.
   static bool hasAttachedPropertyWrapper(VarDecl *vd) {
@@ -1171,7 +1201,7 @@ NO_NEW_INSERTION_POINT(CaptureListScope)
 NO_NEW_INSERTION_POINT(CaseStmtScope)
 NO_NEW_INSERTION_POINT(CaseLabelItemScope)
 NO_NEW_INSERTION_POINT(CaseStmtBodyScope)
-NO_NEW_INSERTION_POINT(ClosureBodyScope)
+NO_NEW_INSERTION_POINT(ClosureParametersScope)
 NO_NEW_INSERTION_POINT(DefaultArgumentInitializerScope)
 NO_NEW_INSERTION_POINT(DoStmtScope)
 NO_NEW_INSERTION_POINT(DoCatchStmtScope)
@@ -1183,10 +1213,8 @@ NO_NEW_INSERTION_POINT(SubscriptDeclScope)
 NO_NEW_INSERTION_POINT(SwitchStmtScope)
 NO_NEW_INSERTION_POINT(VarDeclScope)
 NO_NEW_INSERTION_POINT(WhileStmtScope)
-NO_NEW_INSERTION_POINT(WholeClosureScope)
 
 NO_EXPANSION(GenericParamScope)
-NO_EXPANSION(ClosureParametersScope)
 NO_EXPANSION(SpecializeAttributeScope)
 NO_EXPANSION(DifferentiableAttributeScope)
 NO_EXPANSION(ConditionalClausePatternUseScope)
@@ -1525,35 +1553,15 @@ void SubscriptDeclScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
   scopeCreator.addChildrenForAllLocalizableAccessorsInSourceOrder(sub, params);
 }
 
-void WholeClosureScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
-    ScopeCreator &scopeCreator) {
-  if (auto *cl = captureList.getPtrOrNull())
-    scopeCreator.ensureUniqueThenConstructExpandAndInsert<CaptureListScope>(
-        this, cl);
-  ASTScopeImpl *bodyParent = this;
-  if (closureExpr->getInLoc().isValid())
-    bodyParent =
-        scopeCreator
-            .constructExpandAndInsertUncheckable<ClosureParametersScope>(
-                this, closureExpr, captureList);
-  scopeCreator.constructExpandAndInsertUncheckable<ClosureBodyScope>(
-      bodyParent, closureExpr, captureList);
-}
-
 void CaptureListScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
     ScopeCreator &scopeCreator) {
-  // Patterns here are implicit, so need to dig out the intializers
-  for (const CaptureListEntry &captureListEntry : expr->getCaptureList()) {
-    for (unsigned patternEntryIndex = 0;
-         patternEntryIndex < captureListEntry.Init->getNumPatternEntries();
-         ++patternEntryIndex) {
-      Expr *init = captureListEntry.Init->getInit(patternEntryIndex);
-      scopeCreator.addExprToScopeTree(init, this);
-    }
-  }
+  auto *closureExpr = expr->getClosureBody();
+  scopeCreator
+      .ifUniqueConstructExpandAndInsert<ClosureParametersScope>(
+          this, closureExpr);
 }
 
-void ClosureBodyScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
+void ClosureParametersScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
     ScopeCreator &scopeCreator) {
   scopeCreator.addToScopeTree(closureExpr->getBody(), this);
 }
@@ -1733,51 +1741,6 @@ bool ASTScopeImpl::isATypeDeclScope() const {
   return pd && (isa<NominalTypeDecl>(pd) || isa<ExtensionDecl>(pd));
 }
 
-void ScopeCreator::forEachClosureIn(
-    Expr *expr, function_ref<void(NullablePtr<CaptureListExpr>, ClosureExpr *)>
-                    foundClosure) {
-  ASTScopeAssert(expr,
-                 "If looking for closures, must have an expression to search.");
-
-  /// AST walker that finds top-level closures in an expression.
-  class ClosureFinder : public ASTWalker {
-    function_ref<void(NullablePtr<CaptureListExpr>, ClosureExpr *)>
-        foundClosure;
-
-  public:
-    ClosureFinder(
-        function_ref<void(NullablePtr<CaptureListExpr>, ClosureExpr *)>
-            foundClosure)
-        : foundClosure(foundClosure) {}
-
-    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-      if (auto *closure = dyn_cast<ClosureExpr>(E)) {
-        foundClosure(nullptr, closure);
-        return {false, E};
-      }
-      if (auto *capture = dyn_cast<CaptureListExpr>(E)) {
-        foundClosure(capture, capture->getClosureBody());
-        return {false, E};
-      }
-      return {true, E};
-    }
-    std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
-      if (isa<BraceStmt>(S)) { // closures hidden in here
-        return {true, S};
-      }
-      return {false, S};
-    }
-    std::pair<bool, Pattern *> walkToPatternPre(Pattern *P) override {
-      return {false, P};
-    }
-    bool walkToDeclPre(Decl *D) override { return false; }
-    bool walkToTypeReprPre(TypeRepr *T) override { return false; }
-    bool walkToParameterListPre(ParameterList *PL) override { return false; }
-  };
-
-  expr->walk(ClosureFinder(foundClosure));
-}
-
 #pragma mark new operators
 void *ASTScopeImpl::operator new(size_t bytes, const ASTContext &ctx,
                                  unsigned alignment) {
@@ -1839,7 +1802,7 @@ GET_REFERRENT(VarDeclScope, getDecl())
 GET_REFERRENT(GenericParamScope, paramList->getParams()[index])
 GET_REFERRENT(AbstractStmtScope, getStmt())
 GET_REFERRENT(CaptureListScope, getExpr())
-GET_REFERRENT(WholeClosureScope, getExpr())
+GET_REFERRENT(ClosureParametersScope, getExpr())
 GET_REFERRENT(SpecializeAttributeScope, specializeAttr)
 GET_REFERRENT(DifferentiableAttributeScope, differentiableAttr)
 GET_REFERRENT(GenericTypeOrExtensionScope, portion->getReferrentOfScope(this));
@@ -1980,13 +1943,6 @@ bool PatternEntryDeclScope::isCurrentIfWasExpanded() const {
   return getPatternEntry().getNumBoundVariables() == varCountWhenLastExpanded;
 }
 
-void WholeClosureScope::beCurrent() {
-  bodyWhenLastExpanded = closureExpr->getBody();
-}
-bool WholeClosureScope::isCurrentIfWasExpanded() const {
-  return bodyWhenLastExpanded == closureExpr->getBody();
-}
-
 #pragma mark getParentOfASTAncestorScopesToBeRescued
 NullablePtr<ASTScopeImpl>
 ASTScopeImpl::getParentOfASTAncestorScopesToBeRescued() {

lib/AST/ASTScopeLookup.cpp

@@ -424,18 +424,19 @@ bool PatternEntryInitializerScope::lookupLocalsOrMembers(
   return false;
 }
 
-bool ClosureParametersScope::lookupLocalsOrMembers(
+bool CaptureListScope::lookupLocalsOrMembers(
     ArrayRef<const ASTScopeImpl *>, DeclConsumer consumer) const {
-  if (auto *cl = captureList.getPtrOrNull()) {
-    CaptureListExpr *mutableCL =
-        const_cast<CaptureListExpr *>(captureList.get());
-    for (auto &e : mutableCL->getCaptureList()) {
-      if (consumer.consume(
-              {e.Var},
-              DeclVisibilityKind::LocalVariable)) // or FunctionParameter??
-        return true;
-    }
+  for (auto &e : expr->getCaptureList()) {
+    if (consumer.consume(
+            {e.Var},
+            DeclVisibilityKind::LocalVariable)) // or FunctionParameter??
+      return true;
   }
+  return false;
+}
+
+bool ClosureParametersScope::lookupLocalsOrMembers(
+    ArrayRef<const ASTScopeImpl *>, DeclConsumer consumer) const {
   for (auto param : *closureExpr->getParameters())
     if (consumer.consume({param}, DeclVisibilityKind::FunctionParameter))
       return true;
@@ -702,10 +703,6 @@ Optional<bool> ClosureParametersScope::resolveIsCascadingUseForThisScope(
     Optional<bool> isCascadingUse) const {
   return ifUnknownIsCascadingUseAccordingTo(isCascadingUse, closureExpr);
 }
-Optional<bool> ClosureBodyScope::resolveIsCascadingUseForThisScope(
-    Optional<bool> isCascadingUse) const {
-  return ifUnknownIsCascadingUseAccordingTo(isCascadingUse, closureExpr);
-}
 
 Optional<bool> PatternEntryInitializerScope::resolveIsCascadingUseForThisScope(
     Optional<bool> isCascadingUse) const {