[AST] Clean up handling of single-expression closures

include/swift/AST/ASTWalker.h

@@ -216,12 +216,12 @@ class ASTWalker {
   virtual bool shouldWalkIntoLazyInitializers() { return true; }
 
   /// This method configures whether the walker should visit the body of a
-  /// non-single expression closure.
+  /// closure that was checked separately from its enclosing expression.
   ///
   /// For work that is performed for every top-level expression, this should
   /// be overridden to return false, to avoid duplicating work or visiting
   /// bodies of closures that have not yet been type checked.
-  virtual bool shouldWalkIntoNonSingleExpressionClosure(ClosureExpr *) {
+  virtual bool shouldWalkIntoSeparatelyCheckedClosure(ClosureExpr *) {
     return true;
   }
 

include/swift/AST/AnyFunctionRef.h

@@ -89,20 +89,6 @@ class AnyFunctionRef {
     return TheFunction.get<AbstractClosureExpr *>()->getSingleExpressionBody();
   }
 
-  void setSingleExpressionBody(Expr *expr) {
-    if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
-      AFD->setSingleExpressionBody(expr);
-      return;
-    }
-
-    auto ACE = TheFunction.get<AbstractClosureExpr *>();
-    if (auto CE = dyn_cast<ClosureExpr>(ACE)) {
-      CE->setSingleExpressionBody(expr);
-    } else {
-      cast<AutoClosureExpr>(ACE)->setBody(expr);
-    }
-  }
-
   Type getType() const {
     if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>())
       return AFD->getInterfaceType();

include/swift/AST/Expr.h

@@ -3770,10 +3770,7 @@ class ClosureExpr : public AbstractClosureExpr {
   /// the CaptureListExpr which would normally maintain this sort of
   /// information about captured variables), we need to have some way to access
   /// this information directly on the ClosureExpr.
-  ///
-  /// The bit indicates whether this closure has had a function builder
-  /// applied to it.
-  llvm::PointerIntPair<VarDecl *, 1, bool> CapturedSelfDeclAndAppliedBuilder;
+  VarDecl * CapturedSelfDecl;
 
   /// The location of the "throws", if present.
   SourceLoc ThrowsLoc;
@@ -3787,7 +3784,7 @@ class ClosureExpr : public AbstractClosureExpr {
 
   /// The explicitly-specified result type.
   llvm::PointerIntPair<TypeExpr *, 1, bool>
-    ExplicitResultTypeAndEnclosingChecked;
+    ExplicitResultTypeAndSeparatelyChecked;
 
   /// The body of the closure, along with a bit indicating whether it
   /// was originally just a single expression.
@@ -3800,9 +3797,9 @@ class ClosureExpr : public AbstractClosureExpr {
     : AbstractClosureExpr(ExprKind::Closure, Type(), /*Implicit=*/false,
                           discriminator, parent),
       BracketRange(bracketRange),
-      CapturedSelfDeclAndAppliedBuilder(capturedSelfDecl, false),
+      CapturedSelfDecl(capturedSelfDecl),
       ThrowsLoc(throwsLoc), ArrowLoc(arrowLoc), InLoc(inLoc),
-      ExplicitResultTypeAndEnclosingChecked(explicitResultType, false),
+      ExplicitResultTypeAndSeparatelyChecked(explicitResultType, false),
       Body(nullptr) {
     setParameterList(params);
     Bits.ClosureExpr.HasAnonymousClosureVars = false;
@@ -3857,14 +3854,14 @@ class ClosureExpr : public AbstractClosureExpr {
 
   Type getExplicitResultType() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultTypeAndEnclosingChecked.getPointer()
+    return ExplicitResultTypeAndSeparatelyChecked.getPointer()
         ->getInstanceType();
   }
   void setExplicitResultType(Type ty);
 
   TypeRepr *getExplicitResultTypeRepr() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultTypeAndEnclosingChecked.getPointer()
+    return ExplicitResultTypeAndSeparatelyChecked.getPointer()
         ->getTypeRepr();
   }
 
@@ -3894,42 +3891,27 @@ class ClosureExpr : public AbstractClosureExpr {
   /// Only valid when \c hasSingleExpressionBody() is true.
   Expr *getSingleExpressionBody() const;
 
-  /// Set the body for a closure that has a single expression as its
-  /// body.
-  ///
-  /// This routine cannot change whether a closure has a single expression as
-  /// its body; it can only update that expression.
-  void setSingleExpressionBody(Expr *NewBody);
-
   /// Is this a completely empty closure?
   bool hasEmptyBody() const;
 
   /// VarDecl captured by this closure under the literal name \c self , if any.
   VarDecl *getCapturedSelfDecl() const {
-    return CapturedSelfDeclAndAppliedBuilder.getPointer();
+    return CapturedSelfDecl;
   }
 
   /// Whether this closure captures the \c self param in its body in such a
   /// way that implicit \c self is enabled within its body (i.e. \c self is
   /// captured non-weakly).
   bool capturesSelfEnablingImplictSelf() const;
 
-  bool hasAppliedFunctionBuilder() const {
-    return CapturedSelfDeclAndAppliedBuilder.getInt();
-  }
-
-  void setAppliedFunctionBuilder(bool flag = true) {
-    CapturedSelfDeclAndAppliedBuilder.setInt(flag);
-  }
-
-  /// Whether this closure's body was type checked within the enclosing
-  /// context.
-  bool wasTypeCheckedInEnclosingContext() const {
-    return ExplicitResultTypeAndEnclosingChecked.getInt();
+  /// Whether this closure's body was type checked separately from its
+  /// enclosing expression.
+  bool wasSeparatelyTypeChecked() const {
+    return ExplicitResultTypeAndSeparatelyChecked.getInt();
   }
 
-  void setTypeCheckedInEnclosingContext(bool flag = true) {
-    ExplicitResultTypeAndEnclosingChecked.setInt(flag);
+  void setSeparatelyTypeChecked(bool flag = true) {
+    ExplicitResultTypeAndSeparatelyChecked.setInt(flag);
   }
 
   static bool classof(const Expr *E) {

lib/AST/ASTWalker.cpp

@@ -815,21 +815,15 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
         return nullptr;
     }
 
-    // Handle single-expression closures.
-    if (expr->hasSingleExpressionBody()) {
-      if (Expr *body = doIt(expr->getSingleExpressionBody())) {
-        expr->setSingleExpressionBody(body);
-        return expr;
-      }
-      return nullptr;
-    }
-
-    if (!Walker.shouldWalkIntoNonSingleExpressionClosure(expr))
+    // If the closure was separately type checked and we don't want to
+    // visit separately-checked closure bodies, bail out now.
+    if (expr->wasSeparatelyTypeChecked() &&
+        !Walker.shouldWalkIntoSeparatelyCheckedClosure(expr))
       return expr;
 
     // Handle other closures.
     if (BraceStmt *body = cast_or_null<BraceStmt>(doIt(expr->getBody()))) {
-      expr->setBody(body, false);
+      expr->setBody(body, expr->hasSingleExpressionBody());
       return expr;
     }
     return nullptr;

lib/AST/Expr.cpp

@@ -1898,19 +1898,6 @@ Expr *ClosureExpr::getSingleExpressionBody() const {
   return body.get<Expr *>();
 }
 
-void ClosureExpr::setSingleExpressionBody(Expr *NewBody) {
-  assert(hasSingleExpressionBody() && "Not a single-expression body");
-  auto body = getBody()->getFirstElement();
-  if (auto stmt = body.dyn_cast<Stmt *>()) {
-    if (auto braceStmt = dyn_cast<BraceStmt>(stmt))
-      braceStmt->getFirstElement() = NewBody;
-    else
-      cast<ReturnStmt>(stmt)->setResult(NewBody);
-    return;
-  }
-  getBody()->setFirstElement(NewBody);
-}
-
 bool ClosureExpr::hasEmptyBody() const {
   return getBody()->empty();
 }
@@ -1923,7 +1910,7 @@ bool ClosureExpr::capturesSelfEnablingImplictSelf() const {
 
 void ClosureExpr::setExplicitResultType(Type ty) {
   assert(ty && !ty->hasTypeVariable());
-  ExplicitResultTypeAndEnclosingChecked.getPointer()
+  ExplicitResultTypeAndSeparatelyChecked.getPointer()
       ->setType(MetatypeType::get(ty));
 }
 

lib/IDE/Refactoring.cpp

@@ -3659,16 +3659,25 @@ static CallExpr *findTrailingClosureTarget(SourceManager &SM,
            return N.isStmt(StmtKind::Brace) || N.isExpr(ExprKind::Call);
          });
   Finder.resolve();
-  if (Finder.getContexts().empty()
-      || !Finder.getContexts().back().is<Expr*>())
+  auto contexts = Finder.getContexts();
+  if (contexts.empty())
     return nullptr;
-  CallExpr *CE = cast<CallExpr>(Finder.getContexts().back().get<Expr*>());
+
+  // If the innermost context is a statement (which will be a BraceStmt per
+  // the filtering condition above), drop it.
+  if (contexts.back().is<Stmt *>()) {
+    contexts = contexts.drop_back();
+  }
+
+  if (contexts.empty() || !contexts.back().is<Expr*>())
+    return nullptr;
+  CallExpr *CE = cast<CallExpr>(contexts.back().get<Expr*>());
 
   if (CE->hasTrailingClosure())
     // Call expression already has a trailing closure.
     return nullptr;
 
-  // The last arugment is a closure?
+  // The last argument is a closure?
   Expr *Args = CE->getArg();
   if (!Args)
     return nullptr;

lib/Sema/CSApply.cpp

@@ -3859,7 +3859,8 @@ namespace {
       if (expr->isLiteralInit()) {
         auto *literalInit = expr->getSubExpr();
         if (auto *call = dyn_cast<CallExpr>(literalInit)) {
-          call->getFn()->forEachChildExpr([&](Expr *subExpr) -> Expr * {
+          forEachExprInConstraintSystem(call->getFn(),
+                                        [&](Expr *subExpr) -> Expr * {
             auto *TE = dyn_cast<TypeExpr>(subExpr);
             if (!TE)
               return subExpr;
@@ -5788,7 +5789,7 @@ static bool applyTypeToClosureExpr(ConstraintSystem &cs,
     cs.setType(CE, toType);
 
     // If this closure isn't type-checked in its enclosing expression, write
-    // theg type into the ClosureExpr directly here, since the visitor won't.
+    // the type into the ClosureExpr directly here, since the visitor won't.
     if (!shouldTypeCheckInEnclosingExpression(CE))
       CE->setType(toType);
 

lib/Sema/CSClosure.cpp

@@ -84,6 +84,7 @@ class ClosureConstraintGenerator
     if (!expr)
       return;
 
+    // FIXME: Use SolutionApplicationTarget?
     expr = cs.generateConstraints(expr, closure, /*isInputExpression=*/false);
     if (!expr) {
       hadError = true;
@@ -330,8 +331,6 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
 
     fn.setBody(newBody, /*isSingleExpression=*/false);
     if (closure) {
-      closure->setAppliedFunctionBuilder();
-      closure->setTypeCheckedInEnclosingContext();
       solution.setExprTypes(closure);
     }
 
@@ -347,7 +346,6 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
         solution, closure, closureFnType->getResult(), rewriteTarget);
     application.visit(fn.getBody());
 
-    closure->setTypeCheckedInEnclosingContext();
     return SolutionApplicationToFunctionResult::Success;
   }
 

lib/Sema/CSDiagnostics.cpp

@@ -442,7 +442,7 @@ bool MissingConformanceFailure::diagnoseAsError() {
       }
 
       bool hasFix = false;
-      caseExpr->forEachChildExpr([&](Expr *expr) -> Expr * {
+      forEachExprInConstraintSystem(caseExpr, [&](Expr *expr) -> Expr * {
         hasFix |= anchors.count(expr);
         return hasFix ? nullptr : expr;
       });
@@ -3277,7 +3277,7 @@ bool MissingMemberFailure::diagnoseAsError() {
 
     bool hasUnresolvedPattern = false;
     if (auto *E = getAsExpr(anchor)) {
-      const_cast<Expr *>(E)->forEachChildExpr([&](Expr *expr) {
+      forEachExprInConstraintSystem(const_cast<Expr *>(E), [&](Expr *expr) {
         hasUnresolvedPattern |= isa<UnresolvedPatternExpr>(expr);
         return hasUnresolvedPattern ? nullptr : expr;
       });

lib/Sema/CSGen.cpp

@@ -3882,6 +3882,13 @@ namespace {
           }
         }
 
+        // If this is a closure, only walk into its children if they
+        // are type-checked in the context of the enclosing expression.
+        if (auto closure = dyn_cast<ClosureExpr>(expr)) {
+          if (!shouldTypeCheckInEnclosingExpression(closure))
+            return { false, expr };
+        }
+
         // Now, we're ready to walk into sub expressions.
         return {true, expr};
       }
@@ -4021,12 +4028,6 @@ namespace {
 
     /// Ignore declarations.
     bool walkToDeclPre(Decl *decl) override { return false; }
-
-    // Don't walk into statements.  This handles the BraceStmt in
-    // non-single-expr closures, so we don't walk into their body.
-    std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
-      return { false, S };
-    }
   };
 
   class ConstraintWalker : public ASTWalker {

lib/Sema/CSSimplify.cpp

@@ -1548,7 +1548,7 @@ static bool fixMissingArguments(ConstraintSystem &cs, ASTNode anchor,
 
       // Something like `foo { x in }` or `foo { $0 }`
       if (auto *closure = getAsExpr<ClosureExpr>(anchor)) {
-        closure->forEachChildExpr([&](Expr *expr) -> Expr * {
+        forEachExprInConstraintSystem(closure, [&](Expr *expr) -> Expr * {
           if (auto *UDE = dyn_cast<UnresolvedDotExpr>(expr)) {
             if (!isParam(UDE->getBase()))
               return expr;
@@ -9334,7 +9334,7 @@ bool ConstraintSystem::recordFix(ConstraintFix *fix, unsigned impact) {
 
   bool found = false;
   if (auto *expr = getAsExpr(anchor)) {
-    expr->forEachChildExpr([&](Expr *subExpr) -> Expr * {
+    forEachExprInConstraintSystem(expr, [&](Expr *subExpr) -> Expr * {
       found |= anchors.count(subExpr);
       return subExpr;
     });

lib/Sema/ConstraintSystem.cpp

@@ -590,6 +590,8 @@ static void extendDepthMap(
    Expr *expr,
    llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &depthMap) {
   class RecordingTraversal : public ASTWalker {
+    SmallVector<ClosureExpr *, 4> Closures;
+
   public:
     llvm::DenseMap<Expr *, std::pair<unsigned, Expr *>> &DepthMap;
     unsigned Depth = 0;
@@ -601,13 +603,37 @@ static void extendDepthMap(
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
       DepthMap[E] = {Depth, Parent.getAsExpr()};
       ++Depth;
+
+      if (auto CE = dyn_cast<ClosureExpr>(E))
+        Closures.push_back(CE);
+
       return { true, E };
     }
 
     Expr *walkToExprPost(Expr *E) override {
+      if (auto CE = dyn_cast<ClosureExpr>(E)) {
+        assert(Closures.back() == CE);
+        Closures.pop_back();
+      }
+
       --Depth;
       return E;
     }
+
+    std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
+      if (auto RS = dyn_cast<ReturnStmt>(S)) {
+        // For return statements, treat the parent of the return expression
+        // as the closure itself.
+        if (RS->hasResult() && !Closures.empty()) {
+          llvm::SaveAndRestore<ParentTy> SavedParent(Parent, Closures.back());
+          auto E = RS->getResult();
+          E->walk(*this);
+          return { false, S };
+        }
+      }
+
+      return { true, S };
+    }
   };
 
   RecordingTraversal traversal(depthMap);

lib/Sema/ConstraintSystem.h

@@ -5493,6 +5493,12 @@ BraceStmt *applyFunctionBuilderTransform(
 /// of the parameter and result types without looking at the body.
 bool shouldTypeCheckInEnclosingExpression(ClosureExpr *expr);
 
+/// Visit each subexpression that will be part of the constraint system
+/// of the given expression, including those in closure bodies that will be
+/// part of the constraint system.
+void forEachExprInConstraintSystem(
+    Expr *expr, llvm::function_ref<Expr *(Expr *)> callback);
+
 } // end namespace swift
 
 #endif // LLVM_SWIFT_SEMA_CONSTRAINT_SYSTEM_H