[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.

include/swift/Sema/ConstraintLocator.h

@@ -310,6 +310,10 @@ public:
   /// 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 @@ private:
 // 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;
 
+  if (!path.empty()) {
     // Ignore a trailing SyntacticElement path element for the statement.
-  auto path = getPath();
-  if (auto elt = getLastElementAs<LocatorPathElt::SyntacticElement>()) {
+    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))
+      // 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

test/IDE/complete_if_switch_expr.swift

@@ -1,11 +1,13 @@
 // RUN: %empty-directory(%t)
-// RUN: %target-swift-ide-test -batch-code-completion -source-filename %s -filecheck %raw-FileCheck -completion-output-dir %t
+// RUN: %target-swift-ide-test -batch-code-completion -source-filename %s -debug-forbid-typecheck-prefix FORBIDDEN -filecheck %raw-FileCheck -completion-output-dir %t
 
 enum E {
   case e
   case f(Int)
 }
 
+struct SomeError: Error {}
+
 func ifExprDotReturn() -> E {
   if .random() {
     .#^DOT1?check=DOT^#
@@ -113,6 +115,227 @@ func switchExprBinding() -> E {
   return x
 }
 
+// Triggering interface type computation of this will cause an assert to fire,
+// ensuring we don't attempt to type-check any references to it.
+let NOTYPECHECK = FORBIDDEN
+
+func testSkipTypechecking1() -> E {
+  // Make sure we don't try to type-check the first branch, as it doesn't
+  // contribute to the type.
+  if .random() {
+    _ = NOTYPECHECK
+    throw SomeError()
+  } else {
+    .#^DOT11?check=DOT^#
+  }
+}
+
+func testSkipTypechecking2() -> E {
+  // Make sure we don't try to type-check the second branch, as it doesn't
+  // contribute to the type.
+  if .random() {
+    .#^DOT12?check=DOT^#
+  } else {
+    _ = NOTYPECHECK
+    throw SomeError()
+  }
+}
+
+func testSkipTypechecking3() -> E {
+  // Make sure we don't try to type-check the first branch, as it doesn't
+  // contribute to the type.
+  switch Bool.random() {
+  case true:
+    _ = NOTYPECHECK
+    throw SomeError()
+  case false:
+    .#^DOT13?check=DOT^#
+  }
+}
+
+func testSkipTypechecking4() -> E {
+  // Make sure we don't try to type-check the second branch, as it doesn't
+  // contribute to the type.
+  switch Bool.random() {
+  case true:
+    .#^DOT14?check=DOT^#
+  case false:
+    _ = NOTYPECHECK
+    throw SomeError()
+  }
+}
+
+func takesE(_ e: E) {}
+
+func testSkipTypechecking5() throws -> E {
+  // Make sure we only type-check the first branch.
+  if .random() {
+    // We can still skip type-checking this tho.
+    x = NOTYPECHECK
+
+    takesE(.#^DOT15?check=DOT^#)
+    throw SomeError()
+  } else {
+    NOTYPECHECK
+  }
+}
+
+func testSkipTypechecking6() throws -> E {
+  // Make sure we only type-check the first branch.
+  switch Bool.random() {
+  case true:
+    // We can still skip type-checking this tho.
+    x = NOTYPECHECK
+
+    takesE(.#^DOT16?check=DOT^#)
+    throw SomeError()
+  case false:
+    NOTYPECHECK
+  }
+}
+
+enum F {
+  case x
+}
+
+func takesIntOrDouble(_ i: Int) -> E { E.e }
+func takesIntOrDouble(_ i: Double) -> F { F.x }
+
+func testSkipTypechecking7() throws -> E {
+  let fn = {
+    switch Bool.random() {
+    case true:
+      .#^DOT17?check=DOT^#
+    case false:
+      // Make sure we don't end up with an ambiguity here.
+      takesIntOrDouble(0)
+    }
+  }
+  return fn()
+}
+
+func testSkipTypeChecking8() throws -> E {
+  let e: E = if Bool.random() {
+    takesE(.#^DOT18?check=DOT^#)
+    throw NOTYPECHECK
+  } else {
+    NOTYPECHECK
+  }
+  return e
+}
+
+func testSkipTypeChecking8() throws -> E {
+  // Only need to type-check the inner function in this case.
+  let e: E = if Bool.random() {
+    func localFunc() {
+      takesE(.#^DOT19?check=DOT^#)
+    }
+    throw NOTYPECHECK
+  } else {
+    NOTYPECHECK
+  }
+  return e
+}
+
+func takesArgAndClosure<T>(_ x: Int, fn: () -> T) -> T { fatalError() }
+
+func testSkipTypeChecking9() -> E {
+  // We need to type-check everything for this.
+  if Bool.random() {
+    .e
+  } else {
+    takesArgAndClosure(0) {
+      .#^DOT20?check=DOT^#
+    }
+  }
+}
+
+func testSkipTypeChecking10() -> E {
+  // We only need to type-check the inner-most function for this.
+  if Bool.random() {
+    NOTYPECHECK
+  } else {
+    takesArgAndClosure(NOTYPECHECK) {
+      func foo() {
+        takesE(.#^DOT21?check=DOT^#)
+      }
+      return NOTYPECHECK
+    }
+  }
+}
+
+func testSkipTypeChecking11() -> E {
+  // We only need to type-check the inner-most function for this.
+  if Bool.random() {
+    NOTYPECHECK
+  } else {
+    if NOTYPECHECK {
+      func foo() {
+        takesE(.#^DOT22?check=DOT^#)
+      }
+      throw NOTYPECHECK
+    } else {
+      NOTYPECHECK
+    }
+  }
+}
+
+func testSkipTypeChecking12() -> E {
+  // Only need to type-check the condition here.
+  if .#^DOT23?check=BOOL^# {
+    NOTYPECHECK
+  } else {
+    NOTYPECHECK
+  }
+}
+
+// BOOL:     Begin completions
+// BOOL-DAG: Decl[Constructor]/CurrNominal/IsSystem/TypeRelation[Convertible]: init()[#Bool#]; name=init()
+// BOOL-DAG: Decl[StaticMethod]/CurrNominal/Flair[ExprSpecific]/IsSystem/TypeRelation[Convertible]: random()[#Bool#]; name=random()
+
+func testSkipTypeChecking13(_ e: E) -> E {
+  // TODO: Currently this requires type-checking the bodies, but we ought to
+  // be able to skip them. This is also an issue in closures.
+  switch e {
+  case .#^DOT24?check=DOT^#:
+      .e
+  default:
+      .e
+  }
+}
+
+func testSkipTypeChecking14() -> E {
+  if Bool.random() {
+    .#^DOT25?check=DOT^#
+  } else if .random() {
+    let x = NOTYPECHECK
+    throw x
+  } else if .random() {
+    let x = NOTYPECHECK
+    throw x
+  } else {
+    let x = NOTYPECHECK
+    throw x
+  }
+}
+
+func testSkipTypeChecking14() -> E {
+  switch Bool.random() {
+  case true:
+    .#^DOT26?check=DOT^#
+  case _ where Bool.random():
+    let x = NOTYPECHECK
+    throw x
+  case _ where Bool.random():
+    let x = NOTYPECHECK
+    throw x
+  default:
+    let x = NOTYPECHECK
+    throw x
+  }
+}
+
+
 // DOT:     Begin completions, 2 items
 // DOT-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: e[#E#]; name=e
 // DOT-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: f({#Int#})[#E#]; name=f()