[CodeComplete] More efficient skipping for completions in if/switch exprs

Skip type-checking multi-statement branches if the
completion is in a single-expression branch, and
skip type-checking the expression as a whole if
the completion is in a multi-statement branch.

Author: hamishknight

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/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,

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

@@ -2526,10 +2526,25 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
       if (isa<TapExpr>(E))
         return Action::SkipChildren(E);
 
-      // Don't walk into SingleValueStmtExprs, they should be type-checked as
-      // a whole.
-      if (isa<SingleValueStmtExpr>(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