[AST] Add a enum for type checking state of ClosureExpr's body

Author: rintaro

include/swift/AST/Expr.h

@@ -3759,7 +3759,25 @@ class SerializedAbstractClosureExpr : public SerializedLocalDeclContext {
 ///     { [weak c] (a : Int) -> Int in a + c!.getFoo() }
 /// \endcode
 class ClosureExpr : public AbstractClosureExpr {
+public:
+  enum class BodyState {
+    /// The body was parsed, but not ready for type checking because
+    /// the closure parameters haven't been type checked.
+    Parsed,
+
+    /// The type of the closure itself was type checked. But the body has not
+    /// been type checked yet.
+    ReadyForTypeChecking,
+
+    /// The body was typechecked with the enclosing closure.
+    /// i.e. single expression closure or function builder closure.
+    TypeCheckedAtOnce,
+
+    /// The body was type checked separately from the enclosing closure.
+    SeparatelyTypeChecked,
+  };
 
+private:
   /// The range of the brackets of the capture list, if present.
   SourceRange BracketRange;
 
@@ -3783,8 +3801,7 @@ class ClosureExpr : public AbstractClosureExpr {
   SourceLoc InLoc;
 
   /// The explicitly-specified result type.
-  llvm::PointerIntPair<TypeExpr *, 1, bool>
-    ExplicitResultTypeAndSeparatelyChecked;
+  llvm::PointerIntPair<TypeExpr *, 2, BodyState> ExplicitResultTypeAndBodyState;
 
   /// The body of the closure, along with a bit indicating whether it
   /// was originally just a single expression.
@@ -3799,7 +3816,7 @@ class ClosureExpr : public AbstractClosureExpr {
       BracketRange(bracketRange),
       CapturedSelfDecl(capturedSelfDecl),
       ThrowsLoc(throwsLoc), ArrowLoc(arrowLoc), InLoc(inLoc),
-      ExplicitResultTypeAndSeparatelyChecked(explicitResultType, false),
+      ExplicitResultTypeAndBodyState(explicitResultType, BodyState::Parsed),
       Body(nullptr) {
     setParameterList(params);
     Bits.ClosureExpr.HasAnonymousClosureVars = false;
@@ -3854,15 +3871,13 @@ class ClosureExpr : public AbstractClosureExpr {
 
   Type getExplicitResultType() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultTypeAndSeparatelyChecked.getPointer()
-        ->getInstanceType();
+    return ExplicitResultTypeAndBodyState.getPointer()->getInstanceType();
   }
   void setExplicitResultType(Type ty);
 
   TypeRepr *getExplicitResultTypeRepr() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultTypeAndSeparatelyChecked.getPointer()
-        ->getTypeRepr();
+    return ExplicitResultTypeAndBodyState.getPointer()->getTypeRepr();
   }
 
   /// Determine whether the closure has a single expression for its
@@ -3904,14 +3919,20 @@ class ClosureExpr : public AbstractClosureExpr {
   /// captured non-weakly).
   bool capturesSelfEnablingImplictSelf() const;
 
-  /// Whether this closure's body was type checked separately from its
-  /// enclosing expression.
-  bool wasSeparatelyTypeChecked() const {
-    return ExplicitResultTypeAndSeparatelyChecked.getInt();
+
+  /// Get the type checking state of this closure's body.
+  BodyState getBodyState() const {
+    return ExplicitResultTypeAndBodyState.getInt();
+  }
+  void setBodyState(BodyState v) {
+    ExplicitResultTypeAndBodyState.setInt(v);
   }
 
-  void setSeparatelyTypeChecked(bool flag = true) {
-    ExplicitResultTypeAndSeparatelyChecked.setInt(flag);
+  /// Whether this closure's body is/was type checked separately from its
+  /// enclosing expression.
+  bool isSeparatelyTypeChecked() const {
+    return getBodyState() == BodyState::SeparatelyTypeChecked ||
+           getBodyState() == BodyState::ReadyForTypeChecking;
   }
 
   static bool classof(const Expr *E) {

lib/AST/ASTWalker.cpp

@@ -825,7 +825,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
     // If the closure was separately type checked and we don't want to
     // visit separately-checked closure bodies, bail out now.
-    if (expr->wasSeparatelyTypeChecked() &&
+    if (expr->isSeparatelyTypeChecked() &&
         !Walker.shouldWalkIntoSeparatelyCheckedClosure(expr))
       return expr;
 

lib/AST/Expr.cpp

@@ -1975,7 +1975,7 @@ bool ClosureExpr::capturesSelfEnablingImplictSelf() const {
 
 void ClosureExpr::setExplicitResultType(Type ty) {
   assert(ty && !ty->hasTypeVariable());
-  ExplicitResultTypeAndSeparatelyChecked.getPointer()
+  ExplicitResultTypeAndBodyState.getPointer()
       ->setType(MetatypeType::get(ty));
 }
 

lib/Sema/CSClosure.cpp

@@ -345,11 +345,13 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
     ClosureConstraintApplication application(
         solution, closure, closureFnType->getResult(), rewriteTarget);
     application.visit(fn.getBody());
+    closure->setBodyState(ClosureExpr::BodyState::TypeCheckedAtOnce);
 
     return SolutionApplicationToFunctionResult::Success;
   }
 
   // Otherwise, we need to delay type checking of the closure until later.
   solution.setExprTypes(closure);
+  closure->setBodyState(ClosureExpr::BodyState::ReadyForTypeChecking);
   return SolutionApplicationToFunctionResult::Delay;
 }

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->wasSeparatelyTypeChecked();
+        return !closure->isSeparatelyTypeChecked();
       return false;
     }
 
@@ -1497,7 +1497,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->wasSeparatelyTypeChecked();
+        return !closure->isSeparatelyTypeChecked();
       return false;
     }
 

lib/Sema/TypeCheckAvailability.cpp

@@ -1136,7 +1136,7 @@ static void findAvailabilityFixItNodes(SourceRange ReferenceRange,
         if (Expr *ParentExpr = Parent.getAsExpr()) {
           auto *ParentClosure = dyn_cast<ClosureExpr>(ParentExpr);
           if (!ParentClosure ||
-              ParentClosure->wasSeparatelyTypeChecked()) {
+              ParentClosure->isSeparatelyTypeChecked()) {
             return false;
           }
         } else if (auto *ParentStmt = Parent.getAsStmt()) {

lib/Sema/TypeCheckConstraints.cpp

@@ -2008,7 +2008,7 @@ class FunctionSyntacticDiagnosticWalker : public ASTWalker {
     performSyntacticExprDiagnostics(expr, dcStack.back(), /*isExprStmt=*/false);
 
     if (auto closure = dyn_cast<ClosureExpr>(expr)) {
-      if (closure->wasSeparatelyTypeChecked()) {
+      if (closure->isSeparatelyTypeChecked()) {
         dcStack.push_back(closure);
         return {true, expr};
       }
@@ -2019,7 +2019,7 @@ class FunctionSyntacticDiagnosticWalker : public ASTWalker {
 
   Expr *walkToExprPost(Expr *expr) override {
     if (auto closure = dyn_cast<ClosureExpr>(expr)) {
-      if (closure->wasSeparatelyTypeChecked()) {
+      if (closure->isSeparatelyTypeChecked()) {
         assert(dcStack.back() == closure);
         dcStack.pop_back();
       }

lib/Sema/TypeCheckStmt.cpp

@@ -143,7 +143,7 @@ namespace {
         // 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->wasSeparatelyTypeChecked())
+        if (!CE->isSeparatelyTypeChecked())
           CE->getBody()->walk(ContextualizeClosures(CE));
 
         TypeChecker::computeCaptures(CE);
@@ -1976,7 +1976,7 @@ bool TypeChecker::typeCheckClosureBody(ClosureExpr *closure) {
   if (body) {
     closure->setBody(body, closure->hasSingleExpressionBody());
   }
-  closure->setSeparatelyTypeChecked();
+  closure->setBodyState(ClosureExpr::BodyState::SeparatelyTypeChecked);
   return HadError;
 }