[NFC] AST, Sema: Move TypeChecker::findReturnStatements into AnyFunctionRef

Also rename it to `getExplicitReturnStmts` for clarity and have it
take a `SmallVector` out parameter instead as a small optimization and
to discourage use of this new method as an alternative to
`AnyFunctionRef::bodyHasExplicitReturnStmt`.

Author: AnthonyLatsis

include/swift/AST/AnyFunctionRef.h

@@ -160,6 +160,11 @@ class AnyFunctionRef {
   /// `false` otherwise.
   bool bodyHasExplicitReturnStmt() const;
 
+  /// Finds occurrences of explicit `return` statements within the body, if any.
+  ///
+  /// \param results An out container to which the results are added.
+  void getExplicitReturnStmts(SmallVectorImpl<ReturnStmt *> &results) const;
+
   DeclContext *getAsDeclContext() const {
     if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>())
       return AFD;

include/swift/AST/Decl.h

@@ -109,6 +109,7 @@ namespace swift {
   class ProtocolType;
   struct RawComment;
   enum class ResilienceExpansion : unsigned;
+  class ReturnStmt;
   enum class EffectKind : uint8_t;
   enum class PolymorphicEffectKind : uint8_t;
   class TrailingWhereClause;
@@ -7682,6 +7683,11 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// `false` otherwise.
   bool bodyHasExplicitReturnStmt() const;
 
+  /// Finds occurrences of explicit `return` statements within the body, if any.
+  ///
+  /// \param results An out container to which the results are added.
+  void getExplicitReturnStmts(SmallVectorImpl<ReturnStmt *> &results) 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

@@ -71,6 +71,7 @@ namespace swift {
   class KeyPathExpr;
   class CaptureListExpr;
   class ThenStmt;
+  class ReturnStmt;
 
 enum class ExprKind : uint8_t {
 #define EXPR(Id, Parent) Id,
@@ -4058,6 +4059,11 @@ class AbstractClosureExpr : public DeclContext, public Expr {
   /// `false` otherwise.
   bool bodyHasExplicitReturnStmt() const;
 
+  /// Finds occurrences of explicit `return` statements within the body, if any.
+  ///
+  /// \param results An out container to which the results are added.
+  void getExplicitReturnStmts(SmallVectorImpl<ReturnStmt *> &results) const;
+
   ActorIsolation getActorIsolation() const {
     return actorIsolation;
   }

lib/AST/Decl.cpp

@@ -9361,6 +9361,12 @@ bool AbstractFunctionDecl::bodyHasExplicitReturnStmt() const {
       .bodyHasExplicitReturnStmt();
 }
 
+void AbstractFunctionDecl::getExplicitReturnStmts(
+    SmallVectorImpl<ReturnStmt *> &results) const {
+  AnyFunctionRef(const_cast<AbstractFunctionDecl *>(this))
+      .getExplicitReturnStmts(results);
+}
+
 /// 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

@@ -1967,6 +1967,12 @@ bool AbstractClosureExpr::bodyHasExplicitReturnStmt() const {
       .bodyHasExplicitReturnStmt();
 }
 
+void AbstractClosureExpr::getExplicitReturnStmts(
+    SmallVectorImpl<ReturnStmt *> &results) const {
+  AnyFunctionRef(const_cast<AbstractClosureExpr *>(this))
+      .getExplicitReturnStmts(results);
+}
+
 Type AbstractClosureExpr::getResultType(
     llvm::function_ref<Type(Expr *)> getType) const {
   auto *E = const_cast<AbstractClosureExpr *>(this);

lib/Sema/BuilderTransform.cpp

@@ -918,7 +918,11 @@ 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 (auto returnStmts = findReturnStatements(func); !returnStmts.empty()) {
+
+  SmallVector<ReturnStmt *> returnStmts;
+  func->getExplicitReturnStmts(returnStmts);
+
+  if (!returnStmts.empty()) {
     // One or more explicit 'return' statements were encountered, which
     // disables the result builder transform. Warn when we do this.
     ctx.Diags.diagnose(
@@ -1298,19 +1302,16 @@ bool AnyFunctionRef::bodyHasExplicitReturnStmt() const {
                            BraceHasExplicitReturnStmtRequest{body}, false);
 }
 
-std::vector<ReturnStmt *> TypeChecker::findReturnStatements(AnyFunctionRef fn) {
-  if (!fn.bodyHasExplicitReturnStmt()) {
-    return std::vector<ReturnStmt *>();
+void AnyFunctionRef::getExplicitReturnStmts(
+    SmallVectorImpl<ReturnStmt *> &results) const {
+  if (!bodyHasExplicitReturnStmt()) {
+    return;
   }
 
-  std::vector<ReturnStmt *> results;
-
-  walkExplicitReturnStmts(fn.getBody(), [&results](ReturnStmt *RS) {
+  walkExplicitReturnStmts(getBody(), [&results](ReturnStmt *RS) {
     results.push_back(RS);
     return false;
   });
-
-  return results;
 }
 
 ResultBuilderOpSupport TypeChecker::checkBuilderOpSupport(

lib/Sema/CSDiagnostics.cpp

@@ -8862,7 +8862,8 @@ bool ReferenceToInvalidDeclaration::diagnoseAsError() {
 bool InvalidReturnInResultBuilderBody::diagnoseAsError() {
   auto *closure = castToExpr<ClosureExpr>(getAnchor());
 
-  auto returnStmts = TypeChecker::findReturnStatements(closure);
+  SmallVector<ReturnStmt *> returnStmts;
+  closure->getExplicitReturnStmts(returnStmts);
   assert(!returnStmts.empty());
 
   auto loc = returnStmts.front()->getReturnLoc();

lib/Sema/TypeCheckRequestFunctions.cpp

@@ -263,7 +263,7 @@ static Type inferResultBuilderType(ValueDecl *decl)  {
   // the result of inferring for 'x' will be considered when inferring for 'y'
   // either way.
   if (!funcDecl->hasBody() || !dc->getParentSourceFile() ||
-      !TypeChecker::findReturnStatements(funcDecl).empty()) {
+      funcDecl->bodyHasExplicitReturnStmt()) {
     return Type();
   }
 

lib/Sema/TypeChecker.h

@@ -468,9 +468,6 @@ void typeCheckASTNode(ASTNode &node, DeclContext *DC,
 std::optional<BraceStmt *> applyResultBuilderBodyTransform(FuncDecl *func,
                                                            Type builderType);
 
-/// Find the return statements within the body of the given function.
-std::vector<ReturnStmt *> findReturnStatements(AnyFunctionRef fn);
-
 bool typeCheckTapBody(TapExpr *expr, DeclContext *DC);
 
 Type typeCheckParameterDefault(Expr *&defaultValue, DeclContext *DC,