[CS] Avoid skipping SingleValueStmtExpr branch with ReturnStmt for completion

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 #^^#
+}