[CS] Avoid skipping SingleValueStmtExpr branch with ReturnStmt for completion

We still need to solve a branch with a ReturnStmt to avoid leaving
the contextual result type unbound. This isn't currently legal anyway,
so isn't likely to come up often in practice, but make sure we can
still solve.

Author: hamishknight

include/swift/AST/Stmt.h

@@ -225,6 +225,10 @@ class BraceStmt final : public Stmt,
 
   ASTNode findAsyncNode();
 
+  /// Whether the body contains an explicit `return` statement. This computation
+  /// is cached.
+  bool hasExplicitReturnStmt(ASTContext &ctx) const;
+
   /// If this brace contains a single ASTNode, or a \c #if that has a single active
   /// element, returns it. This will always be the last element of the brace.
   /// Otherwise returns \c nullptr.

lib/AST/Stmt.cpp

@@ -302,6 +302,11 @@ ASTNode BraceStmt::findAsyncNode() {
   return asyncFinder.getAsyncNode();
 }
 
+bool BraceStmt::hasExplicitReturnStmt(ASTContext &ctx) const {
+  return evaluateOrDefault(ctx.evaluator,
+                           BraceHasExplicitReturnStmtRequest{this}, false);
+}
+
 static bool hasSingleActiveElement(ArrayRef<ASTNode> elts) {
   return elts.size() == 1;
 }

lib/Sema/BuilderTransform.cpp

@@ -1289,9 +1289,7 @@ bool AnyFunctionRef::bodyHasExplicitReturnStmt() const {
     return false;
   }
 
-  auto &ctx = getAsDeclContext()->getASTContext();
-  return evaluateOrDefault(ctx.evaluator,
-                           BraceHasExplicitReturnStmtRequest{body}, false);
+  return body->hasExplicitReturnStmt(getAsDeclContext()->getASTContext());
 }
 
 void AnyFunctionRef::getExplicitReturnStmts(

lib/Sema/CSSyntacticElement.cpp

@@ -309,10 +309,13 @@ static bool isViableElement(ASTNode element,
       // 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.
+      // SingleValueStmtExpr. We also need to skip if the body has a ReturnStmt,
+      // which isn't something that's currently allowed, but is necessary to
+      // correctly infer the contextual type without leaving it unbound.
       if (isForSingleValueStmtCompletion &&
           !SingleValueStmtExpr::hasResult(braceStmt) &&
-          !cs.containsIDEInspectionTarget(braceStmt)) {
+          !cs.containsIDEInspectionTarget(braceStmt) &&
+          !braceStmt->hasExplicitReturnStmt(cs.getASTContext())) {
         return false;
       }
     }

test/IDE/complete_singlevaluestmt_return.swift

@@ -0,0 +1,17 @@
+// RUN: %batch-code-completion
+
+struct S {
+  var str: String
+}
+
+_ = {
+  let k = if .random() {
+    return ""
+  } else {
+    S()
+  }
+  // Make sure we can still infer 'k' here.
+  return k.#^COMPLETE_ON_SVE_WITH_RET^#
+  // COMPLETE_ON_SVE_WITH_RET: Decl[InstanceVar]/CurrNominal/TypeRelation[Convertible]: str[#String#]; name=str
+}
+

validation-test/IDE/crashers_fixed/7dfeb692d885c8c5.swift

@@ -0,0 +1,9 @@
+// {"kind":"complete","original":"0bd7af1f","signature":"swift::constraints::TypeVarRefCollector::walkToStmtPre(swift::Stmt*)","signatureAssert":"Assertion failed: (result), function getClosureType"}
+// RUN: %target-swift-ide-test -code-completion -batch-code-completion -skip-filecheck -code-completion-diagnostics -source-filename %s
+{
+  let a =
+    if <#expression#> {
+      return
+    }
+  return #^^#
+}