Merge pull request swiftlang#33981 from slavapestov/simplify-astscope-closure-representation

ASTScope: Simplify representation of closures

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 {