[AST] Record whether a closure was type-checked in its enclosing expression

Rather than using various "applied function builder" and "is single
expression body" checks to determine whether a closure was
type-checked in its enclosing expression, record in the closure
expression whether it actually *was* type-checked as part of its
enclosing expression.

Author: DougGregor

include/swift/AST/Expr.h

@@ -3774,7 +3774,7 @@ class ClosureExpr : public AbstractClosureExpr {
   /// The bit indicates whether this closure has had a function builder
   /// applied to it.
   llvm::PointerIntPair<VarDecl *, 1, bool> CapturedSelfDeclAndAppliedBuilder;
-  
+
   /// The location of the "throws", if present.
   SourceLoc ThrowsLoc;
 
@@ -3786,7 +3786,8 @@ class ClosureExpr : public AbstractClosureExpr {
   SourceLoc InLoc;
 
   /// The explicitly-specified result type.
-  TypeExpr *ExplicitResultType;
+  llvm::PointerIntPair<TypeExpr *, 1, bool>
+    ExplicitResultTypeAndEnclosingChecked;
 
   /// The body of the closure, along with a bit indicating whether it
   /// was originally just a single expression.
@@ -3801,7 +3802,8 @@ class ClosureExpr : public AbstractClosureExpr {
       BracketRange(bracketRange),
       CapturedSelfDeclAndAppliedBuilder(capturedSelfDecl, false),
       ThrowsLoc(throwsLoc), ArrowLoc(arrowLoc), InLoc(inLoc),
-      ExplicitResultType(explicitResultType), Body(nullptr) {
+      ExplicitResultTypeAndEnclosingChecked(explicitResultType, false),
+      Body(nullptr) {
     setParameterList(params);
     Bits.ClosureExpr.HasAnonymousClosureVars = false;
   }
@@ -3855,13 +3857,15 @@ class ClosureExpr : public AbstractClosureExpr {
 
   Type getExplicitResultType() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultType->getInstanceType();
+    return ExplicitResultTypeAndEnclosingChecked.getPointer()
+        ->getInstanceType();
   }
   void setExplicitResultType(Type ty);
 
   TypeRepr *getExplicitResultTypeRepr() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultType->getTypeRepr();
+    return ExplicitResultTypeAndEnclosingChecked.getPointer()
+        ->getTypeRepr();
   }
 
   /// Determine whether the closure has a single expression for its
@@ -3918,6 +3922,16 @@ class ClosureExpr : public AbstractClosureExpr {
     CapturedSelfDeclAndAppliedBuilder.setInt(flag);
   }
 
+  /// Whether this closure's body was type checked within the enclosing
+  /// context.
+  bool wasTypeCheckedInEnclosingContext() const {
+    return ExplicitResultTypeAndEnclosingChecked.getInt();
+  }
+
+  void setTypeCheckedInEnclosingContext(bool flag = true) {
+    ExplicitResultTypeAndEnclosingChecked.setInt(flag);
+  }
+
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::Closure;
   }

lib/AST/Expr.cpp

@@ -1923,7 +1923,8 @@ bool ClosureExpr::capturesSelfEnablingImplictSelf() const {
 
 void ClosureExpr::setExplicitResultType(Type ty) {
   assert(ty && !ty->hasTypeVariable());
-  ExplicitResultType->setType(MetatypeType::get(ty));
+  ExplicitResultTypeAndEnclosingChecked.getPointer()
+      ->setType(MetatypeType::get(ty));
 }
 
 FORWARD_SOURCE_LOCS_TO(AutoClosureExpr, Body)

lib/Sema/CSClosure.cpp

@@ -331,6 +331,7 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
     fn.setBody(newBody, /*isSingleExpression=*/false);
     if (closure) {
       closure->setAppliedFunctionBuilder();
+      closure->setTypeCheckedInEnclosingContext();
       solution.setExprTypes(closure);
     }
 
@@ -343,6 +344,7 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
         solution, closure, closureFnType->getResult(), rewriteTarget);
     application.visit(fn.getBody());
 
+    closure->setTypeCheckedInEnclosingContext();
     return SolutionApplicationToFunctionResult::Success;
   }
 

lib/Sema/MiscDiagnostics.cpp

@@ -86,7 +86,7 @@ static void diagSyntacticUseRestrictions(const Expr *E, const DeclContext *DC,
 
     bool walkToDeclPre(Decl *D) override {
       if (auto *closure = dyn_cast<ClosureExpr>(D->getDeclContext()))
-        return closure->hasAppliedFunctionBuilder();
+        return closure->wasTypeCheckedInEnclosingContext();
       return false;
     }
 
@@ -1459,7 +1459,7 @@ static void diagnoseImplicitSelfUseInClosure(const Expr *E,
     // Don't walk into nested decls.
     bool walkToDeclPre(Decl *D) override {
       if (auto *closure = dyn_cast<ClosureExpr>(D->getDeclContext()))
-        return closure->hasAppliedFunctionBuilder();
+        return closure->wasTypeCheckedInEnclosingContext();
       return false;
     }
 

lib/Sema/TypeCheckStmt.cpp

@@ -140,10 +140,10 @@ namespace {
           }
         }
 
-        // If the closure has a single expression body or has had a function
-        // builder applied to it, we need to walk into it with a new sequence.
+        // If the closure was type checked within its enclosing context,
+        // we need to walk into it with a new sequence.
         // Otherwise, it'll have been separately type-checked.
-        if (CE->hasSingleExpressionBody() || CE->hasAppliedFunctionBuilder())
+        if (CE->wasTypeCheckedInEnclosingContext())
           CE->getBody()->walk(ContextualizeClosures(CE));
 
         TypeChecker::computeCaptures(CE);