From f1ad70343035c2715e5434e41312ede4468127b8 Mon Sep 17 00:00:00 2001
From: Hamish Knight <hamish_knight@apple.com>
Date: Sat, 13 Sep 2025 19:14:35 +0100
Subject: [PATCH] [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.
---
 include/swift/AST/Stmt.h                        |  4 ++++
 lib/AST/Stmt.cpp                                |  5 +++++
 lib/Sema/BuilderTransform.cpp                   |  4 +---
 lib/Sema/CSSyntacticElement.cpp                 |  7 +++++--
 test/IDE/complete_singlevaluestmt_return.swift  | 17 +++++++++++++++++
 .../IDE/crashers_fixed/7dfeb692d885c8c5.swift   |  9 +++++++++
 6 files changed, 41 insertions(+), 5 deletions(-)
 create mode 100644 test/IDE/complete_singlevaluestmt_return.swift
 create mode 100644 validation-test/IDE/crashers_fixed/7dfeb692d885c8c5.swift

diff --git a/include/swift/AST/Stmt.h b/include/swift/AST/Stmt.h
index c8f9d4bb3e76d..933b2c5c07411 100644
--- a/include/swift/AST/Stmt.h
+++ b/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.
diff --git a/lib/AST/Stmt.cpp b/lib/AST/Stmt.cpp
index f1460e2383ed6..68e5dcd8e38af 100644
--- a/lib/AST/Stmt.cpp
+++ b/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;
 }
diff --git a/lib/Sema/BuilderTransform.cpp b/lib/Sema/BuilderTransform.cpp
index 2cdd76f45cac0..326fb242476f9 100644
--- a/lib/Sema/BuilderTransform.cpp
+++ b/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(
diff --git a/lib/Sema/CSSyntacticElement.cpp b/lib/Sema/CSSyntacticElement.cpp
index 8981d5004397f..0180f5a19507d 100644
--- a/lib/Sema/CSSyntacticElement.cpp
+++ b/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;
       }
     }
diff --git a/test/IDE/complete_singlevaluestmt_return.swift b/test/IDE/complete_singlevaluestmt_return.swift
new file mode 100644
index 0000000000000..72f93020c5606
--- /dev/null
+++ b/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
+}
+
diff --git a/validation-test/IDE/crashers_fixed/7dfeb692d885c8c5.swift b/validation-test/IDE/crashers_fixed/7dfeb692d885c8c5.swift
new file mode 100644
index 0000000000000..272bcf9e73069
--- /dev/null
+++ b/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 #^^#
+}