[NFC] AST, Sema: Internalize BraceHasReturnRequest evaluation in AnyFunctionRef method

Author: AnthonyLatsis

include/swift/AST/AnyFunctionRef.h

@@ -153,6 +153,13 @@ class AnyFunctionRef {
     llvm_unreachable("autoclosures don't have statement bodies");
   }
 
+  /// Returns a boolean value indicating whether the body, if any, contains
+  /// an explicit `return` statement.
+  ///
+  /// \returns `true` if the body contains an explicit `return` statement,
+  /// `false` otherwise.
+  bool bodyHasExplicitReturnStmt() const;
+
   DeclContext *getAsDeclContext() const {
     if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>())
       return AFD;

include/swift/AST/Decl.h

@@ -7675,6 +7675,13 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// parsed.
   bool hasBody() const;
 
+  /// Returns a boolean value indicating whether the body, if any, contains
+  /// an explicit `return` statement.
+  ///
+  /// \returns `true` if the body contains an explicit `return` statement,
+  /// `false` otherwise.
+  bool bodyHasExplicitReturnStmt() const;
+
   /// Returns true if the text of this function's body can be retrieved either
   /// by extracting the text from the source buffer or reading the inlinable
   /// body from a deserialized swiftmodule.

include/swift/AST/Expr.h

@@ -4051,6 +4051,13 @@ class AbstractClosureExpr : public DeclContext, public Expr {
   /// returns nullptr if the closure doesn't have a body
   BraceStmt *getBody() const;
 
+  /// Returns a boolean value indicating whether the body, if any, contains
+  /// an explicit `return` statement.
+  ///
+  /// \returns `true` if the body contains an explicit `return` statement,
+  /// `false` otherwise.
+  bool bodyHasExplicitReturnStmt() const;
+
   ActorIsolation getActorIsolation() const {
     return actorIsolation;
   }

lib/AST/Decl.cpp

@@ -9356,6 +9356,11 @@ bool AbstractFunctionDecl::hasBody() const {
   }
 }
 
+bool AbstractFunctionDecl::bodyHasExplicitReturnStmt() const {
+  return AnyFunctionRef(const_cast<AbstractFunctionDecl *>(this))
+      .bodyHasExplicitReturnStmt();
+}
+
 /// Expand all preamble macros attached to the given function declaration.
 static std::vector<ASTNode> expandPreamble(AbstractFunctionDecl *func) {
   std::vector<ASTNode> preamble;

lib/AST/Expr.cpp

@@ -1962,6 +1962,11 @@ BraceStmt * AbstractClosureExpr::getBody() const {
   llvm_unreachable("Unknown closure expression");
 }
 
+bool AbstractClosureExpr::bodyHasExplicitReturnStmt() const {
+  return AnyFunctionRef(const_cast<AbstractClosureExpr *>(this))
+      .bodyHasExplicitReturnStmt();
+}
+
 Type AbstractClosureExpr::getResultType(
     llvm::function_ref<Type(Expr *)> getType) const {
   auto *E = const_cast<AbstractClosureExpr *>(this);

lib/Sema/BuilderTransform.cpp

@@ -918,13 +918,9 @@ std::optional<BraceStmt *>
 TypeChecker::applyResultBuilderBodyTransform(FuncDecl *func, Type builderType) {
   // First look for any return statements, and bail if we have any.
   auto &ctx = func->getASTContext();
-  if (evaluateOrDefault(ctx.evaluator, BraceHasReturnRequest{func->getBody()},
-                        false)) {
+  if (auto returnStmts = findReturnStatements(func); !returnStmts.empty()) {
     // One or more explicit 'return' statements were encountered, which
     // disables the result builder transform. Warn when we do this.
-    auto returnStmts = findReturnStatements(func);
-    assert(!returnStmts.empty());
-
     ctx.Diags.diagnose(
         returnStmts.front()->getReturnLoc(),
         diag::result_builder_disabled_by_return_warn, builderType);
@@ -1126,8 +1122,7 @@ ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
   // not apply the result builder transform if it contained an explicit return.
   // To maintain source compatibility, we still need to check for HasReturnStmt.
   // https://github.com/apple/swift/issues/64332.
-  if (evaluateOrDefault(getASTContext().evaluator,
-                        BraceHasReturnRequest{fn.getBody()}, false)) {
+  if (fn.bodyHasExplicitReturnStmt()) {
     // Diagnostic mode means that solver couldn't reach any viable
     // solution, so let's diagnose presence of a `return` statement
     // in the closure body.
@@ -1276,7 +1271,21 @@ bool BraceHasReturnRequest::evaluate(Evaluator &evaluator,
   return !ReturnStmtFinder::find(BS).empty();
 }
 
+bool AnyFunctionRef::bodyHasExplicitReturnStmt() const {
+  auto *body = getBody();
+  if (!body) {
+    return false;
+  }
+
+  auto &ctx = getAsDeclContext()->getASTContext();
+  return evaluateOrDefault(ctx.evaluator, BraceHasReturnRequest{body}, false);
+}
+
 std::vector<ReturnStmt *> TypeChecker::findReturnStatements(AnyFunctionRef fn) {
+  if (!fn.bodyHasExplicitReturnStmt()) {
+    return std::vector<ReturnStmt *>();
+  }
+
   return ReturnStmtFinder::find(fn.getBody());
 }