AST: Allow CaptureListExpr to hold any AbstractClosureExpr.

Semantically, the capture list binding behavior doesn't require the scope
to be an explicit closure, and forming new ClosureExprs during type-checking
is difficult because they have to have preassigned discriminators, unlike
AutoClosureExprs which get discriminators introduced by Sema itself. Allowing
CaptureListExpr to hold an AutoClosureExpr makes it easier to synthesize
CaptureListExprs during type checking, to represent expressions with closure
semantics that evaluate some parts of the expression eagerly in the surrounding
context.

Author: jckarter

include/swift/AST/ASTScope.h

@@ -203,7 +203,7 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
 
 #pragma mark common queries
 public:
-  virtual NullablePtr<ClosureExpr> getClosureIfClosureScope() const;
+  virtual NullablePtr<AbstractClosureExpr> getClosureIfClosureScope() const;
   virtual ASTContext &getASTContext() const;
   virtual NullablePtr<Decl> getDeclIfAny() const { return nullptr; };
   virtual NullablePtr<Stmt> getStmtIfAny() const { return nullptr; };
@@ -1032,17 +1032,17 @@ class CaptureListScope final : public ASTScopeImpl {
 /// For a closure with named parameters, this scope does the local bindings.
 class ClosureParametersScope final : public ASTScopeImpl {
 public:
-  ClosureExpr *const closureExpr;
+  AbstractClosureExpr *const closureExpr;
 
-  ClosureParametersScope(ClosureExpr *closureExpr)
+  ClosureParametersScope(AbstractClosureExpr *closureExpr)
       : closureExpr(closureExpr) {}
   virtual ~ClosureParametersScope() {}
 
   std::string getClassName() const override;
   SourceRange
   getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
 
-  NullablePtr<ClosureExpr> getClosureIfClosureScope() const override {
+  NullablePtr<AbstractClosureExpr> getClosureIfClosureScope() const override {
     return closureExpr;
   }
   NullablePtr<Expr> getExprIfAny() const override { return closureExpr; }
@@ -1541,7 +1541,7 @@ class BraceStmtScope final : public AbstractStmtScope {
   SourceRange
   getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
 
-  NullablePtr<ClosureExpr> parentClosureIfAny() const; // public??
+  NullablePtr<AbstractClosureExpr> parentClosureIfAny() const; // public??
   Stmt *getStmt() const override { return stmt; }
 
 protected:

include/swift/AST/Expr.h

@@ -4144,10 +4144,10 @@ class CaptureListExpr final : public Expr,
     private llvm::TrailingObjects<CaptureListExpr, CaptureListEntry> {
   friend TrailingObjects;
 
-  ClosureExpr *closureBody;
+  AbstractClosureExpr *closureBody;
 
   CaptureListExpr(ArrayRef<CaptureListEntry> captureList,
-                  ClosureExpr *closureBody)
+                  AbstractClosureExpr *closureBody)
     : Expr(ExprKind::CaptureList, /*Implicit=*/false, Type()),
       closureBody(closureBody) {
     Bits.CaptureListExpr.NumCaptures = captureList.size();
@@ -4158,16 +4158,16 @@ class CaptureListExpr final : public Expr,
 public:
   static CaptureListExpr *create(ASTContext &ctx,
                                  ArrayRef<CaptureListEntry> captureList,
-                                 ClosureExpr *closureBody);
+                                 AbstractClosureExpr *closureBody);
 
   ArrayRef<CaptureListEntry> getCaptureList() {
     return {getTrailingObjects<CaptureListEntry>(),
             Bits.CaptureListExpr.NumCaptures};
   }
-  ClosureExpr *getClosureBody() { return closureBody; }
-  const ClosureExpr *getClosureBody() const { return closureBody; }
+  AbstractClosureExpr *getClosureBody() { return closureBody; }
+  const AbstractClosureExpr *getClosureBody() const { return closureBody; }
 
-  void setClosureBody(ClosureExpr *body) { closureBody = body; }
+  void setClosureBody(AbstractClosureExpr *body) { closureBody = body; }
 
   /// This is a bit weird, but the capture list is lexically contained within
   /// the closure, so the ClosureExpr has the full source range.

lib/AST/ASTScope.cpp

@@ -83,11 +83,11 @@ Pattern *AbstractPatternEntryScope::getPattern() const {
   return getPatternEntry().getPattern();
 }
 
-NullablePtr<ClosureExpr> BraceStmtScope::parentClosureIfAny() const {
+NullablePtr<AbstractClosureExpr> BraceStmtScope::parentClosureIfAny() const {
   return !getParent() ? nullptr : getParent().get()->getClosureIfClosureScope();
 }
 
-NullablePtr<ClosureExpr> ASTScopeImpl::getClosureIfClosureScope() const {
+NullablePtr<AbstractClosureExpr> ASTScopeImpl::getClosureIfClosureScope() const {
   return nullptr;
 }
 

lib/AST/ASTScopeSourceRange.cpp

@@ -283,9 +283,12 @@ SourceRange
 BraceStmtScope::getSourceRangeOfThisASTNode(const bool omitAssertions) const {
   // The brace statements that represent closures start their scope at the
   // 'in' keyword, when present.
-  if (auto closure = parentClosureIfAny()) {
-    if (closure.get()->getInLoc().isValid())
-      return SourceRange(closure.get()->getInLoc(), endLoc);
+  if (auto anyClosure = parentClosureIfAny()) {
+    if (auto closure = dyn_cast<ClosureExpr>(parentClosureIfAny().get())) {
+      if (closure->getInLoc().isValid()) {
+        return SourceRange(closure->getInLoc(), endLoc);
+      }
+    }
   }
   return SourceRange(stmt->getStartLoc(), endLoc);
 }
@@ -302,7 +305,10 @@ SourceRange ConditionalClausePatternUseScope::getSourceRangeOfThisASTNode(
 
 SourceRange
 CaptureListScope::getSourceRangeOfThisASTNode(const bool omitAssertions) const {
-  auto *closureExpr = expr->getClosureBody();
+  if (auto autoClosure = dyn_cast<AutoClosureExpr>(expr->getClosureBody())) {
+    return autoClosure->getSourceRange();
+  }
+  auto closureExpr = cast<ClosureExpr>(expr->getClosureBody());
   if (!omitAssertions)
     ASTScopeAssert(closureExpr->getInLoc().isValid(),
                    "We don't create these if no in loc");
@@ -311,10 +317,16 @@ CaptureListScope::getSourceRangeOfThisASTNode(const bool omitAssertions) const {
 
 SourceRange ClosureParametersScope::getSourceRangeOfThisASTNode(
     const bool omitAssertions) const {
-  if (closureExpr->getInLoc().isValid())
-    return SourceRange(closureExpr->getInLoc(), closureExpr->getEndLoc());
+  if (auto autoClosure = dyn_cast<AutoClosureExpr>(closureExpr)) {
+    return autoClosure->getSourceRange();
+  }
+  auto explicitClosureExpr = cast<ClosureExpr>(closureExpr);
+  if (explicitClosureExpr->getInLoc().isValid()) {
+    return SourceRange(explicitClosureExpr->getInLoc(),
+                       explicitClosureExpr->getEndLoc());
+  }
 
-  return closureExpr->getSourceRange();
+  return explicitClosureExpr->getSourceRange();
 }
 
 SourceRange AttachedPropertyWrapperScope::getSourceRangeOfThisASTNode(

lib/AST/ASTWalker.cpp

@@ -848,8 +848,8 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       }
     }
 
-    ClosureExpr *body = expr->getClosureBody();
-    if ((body = cast_or_null<ClosureExpr>(doIt(body))))
+    AbstractClosureExpr *body = expr->getClosureBody();
+    if ((body = cast_or_null<AbstractClosureExpr>(doIt(body))))
       expr->setClosureBody(body);
     else
       return nullptr;

lib/AST/Expr.cpp

@@ -1302,7 +1302,7 @@ bool CaptureListEntry::isSimpleSelfCapture() const {
 
 CaptureListExpr *CaptureListExpr::create(ASTContext &ctx,
                                          ArrayRef<CaptureListEntry> captureList,
-                                         ClosureExpr *closureBody) {
+                                         AbstractClosureExpr *closureBody) {
   auto size = totalSizeToAlloc<CaptureListEntry>(captureList.size());
   auto mem = ctx.Allocate(size, alignof(CaptureListExpr));
   auto *expr = ::new(mem) CaptureListExpr(captureList, closureBody);

lib/IDE/Formatting.cpp

@@ -140,7 +140,7 @@ static ClosureExpr *findTrailingClosureFromArgument(Expr *arg) {
   if (auto TC = dyn_cast_or_null<ClosureExpr>(arg))
     return TC;
   if (auto TCL = dyn_cast_or_null<CaptureListExpr>(arg))
-    return TCL->getClosureBody();
+    return dyn_cast_or_null<ClosureExpr>(TCL->getClosureBody());
   return nullptr;
 }
 
@@ -2488,7 +2488,7 @@ class FormatWalker : public ASTWalker {
   Optional<IndentContext>
   getIndentContextFrom(CaptureListExpr *CL,
                        SourceLoc ContextLoc = SourceLoc()) {
-    ClosureExpr *CE = CL->getClosureBody();
+    AbstractClosureExpr *CE = CL->getClosureBody();
     BraceStmt *BS = CE->getBody();
     if (!CE || !BS)
       return None;
@@ -2503,8 +2503,14 @@ class FormatWalker : public ASTWalker {
   }
 
   Optional<IndentContext>
-  getIndentContextFrom(ClosureExpr *CE, SourceLoc ContextLoc = SourceLoc(),
+  getIndentContextFrom(AbstractClosureExpr *ACE, SourceLoc ContextLoc = SourceLoc(),
                        CaptureListExpr *ParentCapture = nullptr) {
+    // Explicit capture lists should always have an explicit ClosureExpr as
+    // their subexpression.
+    auto CE = dyn_cast<ClosureExpr>(ACE);
+    if (!CE) {
+      return None;
+    }
     BraceStmt *BS = CE->getBody();
     if (!BS)
       return None;

lib/IDE/Refactoring.cpp

@@ -7315,7 +7315,7 @@ class AsyncConverter : private SourceEntityWalker {
     if (auto Closure = dyn_cast<ClosureExpr>(E)) {
       return Closure;
     } else if (auto CaptureList = dyn_cast<CaptureListExpr>(E)) {
-      return CaptureList->getClosureBody();
+      return dyn_cast<ClosureExpr>(CaptureList->getClosureBody());
     } else {
       return nullptr;
     }