[Diagnostics] Diagnose use of one or more return statements in a result builder body

Author: xedin

lib/Sema/CSDiagnostics.cpp

@@ -7052,3 +7052,22 @@ bool ReferenceToInvalidDeclaration::diagnoseAsError() {
   emitDiagnosticAt(decl, diag::decl_declared_here, decl->getName());
   return true;
 }
+
+bool InvalidReturnInResultBuilderBody::diagnoseAsError() {
+  auto *closure = castToExpr<ClosureExpr>(getAnchor());
+
+  auto returnStmts = TypeChecker::findReturnStatements(closure);
+  assert(!returnStmts.empty());
+
+  auto loc = returnStmts.front()->getReturnLoc();
+  emitDiagnosticAt(loc, diag::result_builder_disabled_by_return, BuilderType);
+
+  // Note that one can remove all of the return statements.
+  {
+    auto diag = emitDiagnosticAt(loc, diag::result_builder_remove_returns);
+    for (auto returnStmt : returnStmts)
+      diag.fixItRemove(returnStmt->getReturnLoc());
+  }
+
+  return true;
+}

lib/Sema/CSDiagnostics.h

@@ -2308,6 +2308,26 @@ class ReferenceToInvalidDeclaration final : public FailureDiagnostic {
   bool diagnoseAsError() override;
 };
 
+/// Diagnose use of `return` statements in a body of a result builder.
+///
+/// \code
+/// struct S : Builder {
+///   var foo: some Builder {
+///     return EmptyBuilder()
+///   }
+/// }
+/// \endcode
+class InvalidReturnInResultBuilderBody final : public FailureDiagnostic {
+  Type BuilderType;
+
+public:
+  InvalidReturnInResultBuilderBody(const Solution &solution, Type builderTy,
+                                   ConstraintLocator *locator)
+      : FailureDiagnostic(solution, locator), BuilderType(builderTy) {}
+
+  bool diagnoseAsError() override;
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSFix.cpp

@@ -1647,7 +1647,8 @@ AllowRefToInvalidDecl::create(ConstraintSystem &cs,
 
 bool IgnoreResultBuilderWithReturnStmts::diagnose(const Solution &solution,
                                                   bool asNote) const {
-  return false;
+  InvalidReturnInResultBuilderBody failure(solution, BuilderType, getLocator());
+  return failure.diagnose(asNote);
 }
 
 IgnoreResultBuilderWithReturnStmts *