[CodeComplete] Properly handle if/switch expressions

include/swift/AST/TypeCheckRequests.h

@@ -3868,6 +3868,30 @@ class PreCheckReturnStmtRequest
   bool isCached() const { return true; }
 };
 
+/// Performs some pre-checking of a function body, including inserting and
+/// removing implict returns where needed. This request is currently
+/// side-effectful, as it will mutate the body in-place.
+///
+/// Note this request is currently uncached as:
+/// - The result will ultimately be cached by TypeCheckFunctionBodyRequest.
+/// - When re-parsing function bodies in SourceKit, we can set a previously
+///   type-checked function body back to being parsed, so caching the result of
+///   this request would result in incorrectly attempting to restore the
+///   previous body. We either need to eliminate the mutation of the AST, or
+///   implement request cache invalidation through request dependencies.
+class PreCheckFunctionBodyRequest
+    : public SimpleRequest<PreCheckFunctionBodyRequest,
+                           BraceStmt *(AbstractFunctionDecl *),
+                           RequestFlags::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  BraceStmt *evaluate(Evaluator &evaluator, AbstractFunctionDecl *AFD) const;
+};
+
 /// The result of the query for whether a statement can produce a single value.
 class IsSingleValueStmtResult {
 public:

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -325,6 +325,9 @@ SWIFT_REQUEST(TypeChecker, SynthesizeAccessorRequest,
               SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TangentStoredPropertyRequest,
               llvm::Expected<VarDecl *>(VarDecl *, CanType), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, PreCheckFunctionBodyRequest,
+              BraceStmt *(AbstractFunctionDecl *),
+              Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyRequest,
               BraceStmt *(AbstractFunctionDecl *), SeparatelyCached,
               NoLocationInfo)

include/swift/Sema/ConstraintLocator.h

@@ -310,6 +310,10 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// SingleValueStmtExpr.
   bool isForSingleValueStmtConjunction() const;
 
+  /// Whether this locator identifies a conjunction for the branches of a
+  /// SingleValueStmtExpr, or a conjunction for one of the BraceStmts itself.
+  bool isForSingleValueStmtConjunctionOrBrace() const;
+
   /// Whether this locator identifies a conversion for a SingleValueStmtExpr
   /// branch, and if so, the kind of branch.
   llvm::Optional<SingleValueStmtBranchKind> isForSingleValueStmtBranch() const;

lib/IDE/TypeCheckCompletionCallback.cpp

@@ -52,8 +52,6 @@ Type swift::ide::getTypeForCompletion(const constraints::Solution &S,
     return nullptr;
   }
 
-  auto &CS = S.getConstraintSystem();
-
   Type Result;
 
   if (isExpr<CodeCompletionExpr>(Node)) {
@@ -62,9 +60,9 @@ Type swift::ide::getTypeForCompletion(const constraints::Solution &S,
     Result = S.getResolvedType(Node);
   }
 
-  if (!Result || Result->is<UnresolvedType>()) {
-    Result = CS.getContextualType(Node, /*forConstraint=*/false);
-  }
+  if (!Result || Result->is<UnresolvedType>())
+    Result = S.getContextualType(Node);
+
   if (Result && Result->is<UnresolvedType>()) {
     Result = Type();
   }

lib/Sema/CSStep.cpp

@@ -883,11 +883,15 @@ bool ConjunctionStep::attempt(const ConjunctionElement &element) {
   assert(!ModifiedOptions.has_value() &&
          "Previously modified options should have been restored in resume");
   if (CS.isForCodeCompletion() &&
-      !element.mightContainCodeCompletionToken(CS)) {
+      !element.mightContainCodeCompletionToken(CS) &&
+      !getLocator()->isForSingleValueStmtConjunctionOrBrace()) {
     ModifiedOptions.emplace(CS.Options);
     // If we know that this conjunction element doesn't contain the code
     // completion token, type check it in normal mode without any special
     // behavior that is intended for the code completion token.
+    // Avoid doing this for SingleValueStmtExprs, because we can more eagerly
+    // prune branches in that case, which requires us to detect the code
+    // completion option while solving the conjunction.
     CS.Options -= ConstraintSystemFlags::ForCodeCompletion;
   }
 

lib/Sema/CSSyntacticElement.cpp

@@ -242,7 +242,9 @@ class TypeVariableRefFinder : public ASTWalker {
 // MARK: Constraint generation
 
 /// Check whether it makes sense to convert this element into a constraint.
-static bool isViableElement(ASTNode element) {
+static bool isViableElement(ASTNode element,
+                            bool isForSingleValueStmtCompletion,
+                            ConstraintSystem &cs) {
   if (auto *decl = element.dyn_cast<Decl *>()) {
     // - Ignore variable declarations, they are handled by pattern bindings;
     // - Ignore #if, the chosen children should appear in the
@@ -255,10 +257,21 @@ static bool isViableElement(ASTNode element) {
   }
 
   if (auto *stmt = element.dyn_cast<Stmt *>()) {
-    // Empty brace statements are now viable because they do not require
-    // inference.
     if (auto *braceStmt = dyn_cast<BraceStmt>(stmt)) {
-      return braceStmt->getNumElements() > 0;
+      // Empty brace statements are not viable because they do not require
+      // inference.
+      if (braceStmt->empty())
+        return false;
+
+      // Skip if we're doing completion for a SingleValueStmtExpr, and have a
+      // brace that doesn't involve a single expression, and doesn't have a
+      // code completion token, as it won't contribute to the type of the
+      // SingleValueStmtExpr.
+      if (isForSingleValueStmtCompletion &&
+          !braceStmt->getSingleActiveExpression() &&
+          !cs.containsIDEInspectionTarget(braceStmt)) {
+        return false;
+      }
     }
   }
 
@@ -324,13 +337,19 @@ static void createConjunction(ConstraintSystem &cs,
 
   VarRefCollector paramCollector(cs);
 
+  // Whether we're doing completion, and the conjunction is for a
+  // SingleValueStmtExpr, or one of its braces.
+  const auto isForSingleValueStmtCompletion =
+      cs.isForCodeCompletion() &&
+      locator->isForSingleValueStmtConjunctionOrBrace();
+
   for (const auto &entry : elements) {
     ASTNode element = std::get<0>(entry);
     ContextualTypeInfo context = std::get<1>(entry);
     bool isDiscarded = std::get<2>(entry);
     ConstraintLocator *elementLoc = std::get<3>(entry);
 
-    if (!isViableElement(element))
+    if (!isViableElement(element, isForSingleValueStmtCompletion, cs))
       continue;
 
     // If this conjunction going to represent a body of a closure,
@@ -1122,9 +1141,13 @@ class SyntacticElementConstraintGenerator
 
     for (auto element : braceStmt->getElements()) {
       if (cs.isForCodeCompletion() &&
-          !cs.containsIDEInspectionTarget(element)) {
+          !cs.containsIDEInspectionTarget(element) &&
+          !(braceStmt->getSingleActiveExpression() &&
+            locator->isForSingleValueStmtConjunctionOrBrace())) {
         // To improve performance, skip type checking elements that can't
-        // influence the code completion token.
+        // influence the code completion token. Note we don't do this for
+        // single expression SingleValueStmtExpr branches, as they're needed to
+        // infer the type.
         if (element.is<Stmt *>() && !element.isStmt(StmtKind::Guard) && !element.isStmt(StmtKind::Return)) {
           // Statements can't influence the expresion that contains the code
           // completion token.

lib/Sema/ConstraintLocator.cpp

@@ -634,22 +634,52 @@ bool ConstraintLocator::isForMacroExpansion() const {
   return directlyAt<MacroExpansionExpr>();
 }
 
-bool ConstraintLocator::isForSingleValueStmtConjunction() const {
-  auto *SVE = getAsExpr<SingleValueStmtExpr>(getAnchor());
+static bool isForSingleValueStmtConjunction(ASTNode anchor,
+                                            ArrayRef<LocatorPathElt> path) {
+  auto *SVE = getAsExpr<SingleValueStmtExpr>(anchor);
   if (!SVE)
     return false;
 
-  // Ignore a trailing SyntacticElement path element for the statement.
-  auto path = getPath();
-  if (auto elt = getLastElementAs<LocatorPathElt::SyntacticElement>()) {
-    if (elt->getElement() == ASTNode(SVE->getStmt()))
-      path = path.drop_back();
+  if (!path.empty()) {
+    // Ignore a trailing SyntacticElement path element for the statement.
+    if (auto elt = path.back().getAs<LocatorPathElt::SyntacticElement>()) {
+      if (elt->getElement() == ASTNode(SVE->getStmt()))
+        path = path.drop_back();
+    }
   }
 
   // Other than the trailing SyntaticElement, we must be at the anchor.
   return path.empty();
 }
 
+bool ConstraintLocator::isForSingleValueStmtConjunction() const {
+  return ::isForSingleValueStmtConjunction(getAnchor(), getPath());
+}
+
+bool ConstraintLocator::isForSingleValueStmtConjunctionOrBrace() const {
+  if (!isExpr<SingleValueStmtExpr>(getAnchor()))
+    return false;
+
+  auto path = getPath();
+  while (!path.empty()) {
+    // Ignore a trailing TernaryBranch locator, which is used for if statements.
+    if (path.back().is<LocatorPathElt::TernaryBranch>()) {
+      path = path.drop_back();
+      continue;
+    }
+
+    // Ignore a SyntaticElement path element for a case statement of a switch.
+    if (auto elt = path.back().getAs<LocatorPathElt::SyntacticElement>()) {
+      if (elt->getElement().isStmt(StmtKind::Case)) {
+        path = path.drop_back();
+        continue;
+      }
+    }
+    break;
+  }
+  return ::isForSingleValueStmtConjunction(getAnchor(), path);
+}
+
 llvm::Optional<SingleValueStmtBranchKind>
 ConstraintLocator::isForSingleValueStmtBranch() const {
   // Ignore a trailing ContextualType path element.

lib/Sema/TypeCheckStmt.cpp

@@ -2396,6 +2396,13 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
         return false;
       }
     }
+    if (auto *AFD = dyn_cast<AbstractFunctionDecl>(DC)) {
+      if (AFD->hasBody() && !AFD->isBodyTypeChecked()) {
+        // Pre-check the function body if needed.
+        (void)evaluateOrDefault(evaluator, PreCheckFunctionBodyRequest{AFD},
+                                nullptr);
+      }
+    }
   }
 
   // Find innermost ASTNode at Loc from DC. Results the reference to the found
@@ -2519,6 +2526,26 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
       if (isa<TapExpr>(E))
         return Action::SkipChildren(E);
 
+      // If the location is within a single-expression branch of a
+      // SingleValueStmtExpr, walk it directly rather than as part of the brace.
+      // This ensures we type-check it a part of the whole expression, unless it
+      // has an inner closure or SVE, in which case we can still pick up a
+      // better node to type-check.
+      if (auto *SVE = dyn_cast<SingleValueStmtExpr>(E)) {
+        SmallVector<Expr *> scratch;
+        for (auto *branch : SVE->getSingleExprBranches(scratch)) {
+          auto branchCharRange = Lexer::getCharSourceRangeFromSourceRange(
+            SM, branch->getSourceRange());
+          if (branchCharRange.contains(Loc)) {
+            if (!branch->walk(*this))
+              return Action::Stop();
+
+            return Action::SkipChildren(E);
+          }
+        }
+        return Action::Continue(E);
+      }
+
       if (auto closure = dyn_cast<ClosureExpr>(E)) {
         // NOTE: When a client wants to type check a closure signature, it
         // requests with closure's 'getLoc()' location.
@@ -2599,11 +2626,6 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
         // thus the transform couldn't be applied. Perform code completion
         // pretending there was no result builder to recover.
       }
-    } else if (func->hasSingleExpressionBody() &&
-                func->getResultInterfaceType()->isVoid()) {
-       // The function returns void.  We don't need an explicit return, no matter
-       // what the type of the expression is.  Take the inserted return back out.
-      func->getBody()->setLastElement(func->getSingleExpressionBody());
     }
   }
 
@@ -2641,7 +2663,60 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
 }
 
 BraceStmt *
-TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
+PreCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
+                                      AbstractFunctionDecl *AFD) const {
+  auto &ctx = AFD->getASTContext();
+  auto *body = AFD->getBody();
+  assert(body && "Expected body");
+  assert(!AFD->isBodyTypeChecked() && "Body already type-checked?");
+
+  if (auto *func = dyn_cast<FuncDecl>(AFD)) {
+    // Don't apply this pre-checking to functions with result builders.
+    if (getResultBuilderType(func))
+      return body;
+
+    if (func->hasSingleExpressionBody() &&
+        func->getResultInterfaceType()->isVoid()) {
+      // The function returns void.  We don't need an explicit return, no
+      // matter what the type of the expression is. Take the inserted return
+      // back out.
+      body->setLastElement(func->getSingleExpressionBody());
+    }
+    // If there is a single statement in the body that can be turned into a
+    // single expression return, do so now.
+    if (!func->getResultInterfaceType()->isVoid()) {
+      if (auto *S = body->getSingleActiveStatement()) {
+        if (S->mayProduceSingleValue(evaluator)) {
+          auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
+              ctx, S, /*DC*/ func, /*mustBeExpr*/ false);
+          auto *RS = new (ctx) ReturnStmt(SourceLoc(), SVE);
+          body->setLastElement(RS);
+          func->setHasSingleExpressionBody();
+          func->setSingleExpressionBody(SVE);
+        }
+      }
+    }
+  }
+
+  if (auto *ctor = dyn_cast<ConstructorDecl>(AFD)) {
+    if (body->empty() || !isKnownEndOfConstructor(body->getLastElement())) {
+      // For constructors, we make sure that the body ends with a "return"
+      // stmt, which we either implicitly synthesize, or the user can write.
+      // This simplifies SILGen.
+      SmallVector<ASTNode, 8> Elts(body->getElements().begin(),
+                                   body->getElements().end());
+      Elts.push_back(new (ctx) ReturnStmt(body->getRBraceLoc(),
+                                          /*value*/ nullptr,
+                                          /*implicit*/ true));
+      body = BraceStmt::create(ctx, body->getLBraceLoc(), Elts,
+                               body->getRBraceLoc(), body->isImplicit());
+    }
+  }
+  return body;
+}
+
+BraceStmt *
+TypeCheckFunctionBodyRequest::evaluate(Evaluator &eval,
                                        AbstractFunctionDecl *AFD) const {
   ASTContext &ctx = AFD->getASTContext();
 
@@ -2672,6 +2747,12 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
                              range.End);
   };
 
+  // First do a pre-check of the body.
+  body = evaluateOrDefault(eval, PreCheckFunctionBodyRequest{AFD}, nullptr);
+  assert(body);
+
+  // Then apply a result builder if we have one, which if successful will
+  // produce a type-checked body.
   bool alreadyTypeChecked = false;
   if (auto *func = dyn_cast<FuncDecl>(AFD)) {
     if (Type builderType = getResultBuilderType(func)) {
@@ -2686,42 +2767,6 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
 
         body->walk(ContextualizeClosuresAndMacros(AFD));
       }
-    } else {
-      if (func->hasSingleExpressionBody() &&
-          func->getResultInterfaceType()->isVoid()) {
-        // The function returns void.  We don't need an explicit return, no
-        // matter what the type of the expression is. Take the inserted return
-        // back out.
-        body->setLastElement(func->getSingleExpressionBody());
-      }
-      // If there is a single statement in the body that can be turned into a
-      // single expression return, do so now.
-      if (!func->getResultInterfaceType()->isVoid()) {
-        if (auto *S = body->getSingleActiveStatement()) {
-          if (S->mayProduceSingleValue(evaluator)) {
-            auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
-                ctx, S, /*DC*/ func, /*mustBeExpr*/ false);
-            auto *RS = new (ctx) ReturnStmt(SourceLoc(), SVE);
-            body->setLastElement(RS);
-            func->setHasSingleExpressionBody();
-            func->setSingleExpressionBody(SVE);
-          }
-        }
-      }
-    }
-  } else if (auto *ctor = dyn_cast<ConstructorDecl>(AFD)) {
-    if (body->empty() ||
-        !isKnownEndOfConstructor(body->getLastElement())) {
-      // For constructors, we make sure that the body ends with a "return" stmt,
-      // which we either implicitly synthesize, or the user can write.  This
-      // simplifies SILGen.
-      SmallVector<ASTNode, 8> Elts(body->getElements().begin(),
-                                   body->getElements().end());
-      Elts.push_back(new (ctx) ReturnStmt(body->getRBraceLoc(),
-                                          /*value*/nullptr,
-                                          /*implicit*/true));
-      body = BraceStmt::create(ctx, body->getLBraceLoc(), Elts,
-                               body->getRBraceLoc(), body->isImplicit());
     }
   }