Merge pull request #34810 from xedin/rdar-66709164

[ConstraintSystem] Detect and diagnose extraneous return(s) in result build body

Author: xedin

include/swift/AST/DiagnosticsSema.def

@@ -5482,7 +5482,10 @@ NOTE(previous_result_builder_here, none,
      "previous result builder specified here", ())
 ERROR(result_builder_arguments, none,
       "result builder attributes cannot have arguments", ())
-WARNING(result_builder_disabled_by_return, none,
+ERROR(result_builder_disabled_by_return, none,
+      "cannot use explicit 'return' statement in the body of result builder %0",
+      (Type))
+WARNING(result_builder_disabled_by_return_warn, none,
         "application of result builder %0 disabled by explicit 'return' "
         "statement", (Type))
 NOTE(result_builder_remove_attr, none,

include/swift/Sema/CSFix.h

@@ -2028,7 +2028,8 @@ class AllowKeyPathWithoutComponents final : public ConstraintFix {
                                                ConstraintLocator *locator);
 };
 
-class IgnoreInvalidResultBuilderBody final : public ConstraintFix {
+class IgnoreInvalidResultBuilderBody : public ConstraintFix {
+protected:
   enum class ErrorInPhase {
     PreCheck,
     ConstraintGeneration,
@@ -2067,6 +2068,22 @@ class IgnoreInvalidResultBuilderBody final : public ConstraintFix {
   create(ConstraintSystem &cs, ErrorInPhase phase, ConstraintLocator *locator);
 };
 
+class IgnoreResultBuilderWithReturnStmts final
+    : public IgnoreInvalidResultBuilderBody {
+  Type BuilderType;
+
+  IgnoreResultBuilderWithReturnStmts(ConstraintSystem &cs, Type builderTy,
+                                     ConstraintLocator *locator)
+      : IgnoreInvalidResultBuilderBody(cs, ErrorInPhase::PreCheck, locator),
+        BuilderType(builderTy) {}
+
+public:
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  static IgnoreResultBuilderWithReturnStmts *
+  create(ConstraintSystem &cs, Type builderTy, ConstraintLocator *locator);
+};
+
 class SpecifyContextualTypeForNil final : public ConstraintFix {
   SpecifyContextualTypeForNil(ConstraintSystem &cs,
                               ConstraintLocator *locator)

lib/Sema/BuilderTransform.cpp

@@ -1554,7 +1554,7 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
 
     ctx.Diags.diagnose(
         returnStmts.front()->getReturnLoc(),
-        diag::result_builder_disabled_by_return, builderType);
+        diag::result_builder_disabled_by_return_warn, builderType);
 
     // Note that one can remove the result builder attribute.
     auto attr = func->getAttachedResultBuilder();
@@ -1671,7 +1671,7 @@ ConstraintSystem::matchResultBuilder(
   if (InvalidResultBuilderBodies.count(fn)) {
     (void)recordFix(
         IgnoreInvalidResultBuilderBody::duringConstraintGeneration(
-            *this, getConstraintLocator(fn.getBody())));
+            *this, getConstraintLocator(fn.getAbstractClosureExpr())));
     return getTypeMatchSuccess();
   }
 
@@ -1690,13 +1690,25 @@ ConstraintSystem::matchResultBuilder(
       return getTypeMatchFailure(locator);
 
     if (recordFix(IgnoreInvalidResultBuilderBody::duringPreCheck(
-            *this, getConstraintLocator(fn.getBody()))))
+            *this, getConstraintLocator(fn.getAbstractClosureExpr()))))
       return getTypeMatchFailure(locator);
 
     return getTypeMatchSuccess();
   }
 
   case ResultBuilderBodyPreCheck::HasReturnStmt:
+    // Diagnostic mode means that solver couldn't reach any viable
+    // solution, so let's diagnose presence of a `return` statement
+    // in the closure body.
+    if (shouldAttemptFixes()) {
+      if (recordFix(IgnoreResultBuilderWithReturnStmts::create(
+              *this, builderType,
+              getConstraintLocator(fn.getAbstractClosureExpr()))))
+        return getTypeMatchFailure(locator);
+
+      return getTypeMatchSuccess();
+    }
+
     // If the body has a return statement, suppress the transform but
     // continue solving the constraint system.
     return None;
@@ -1744,7 +1756,7 @@ ConstraintSystem::matchResultBuilder(
 
       if (recordFix(
               IgnoreInvalidResultBuilderBody::duringConstraintGeneration(
-                  *this, getConstraintLocator(fn.getBody()))))
+                  *this, getConstraintLocator(fn.getAbstractClosureExpr()))))
         return getTypeMatchFailure(locator);
 
       return getTypeMatchSuccess();

lib/Sema/CSBindings.cpp

@@ -1454,7 +1454,7 @@ TypeVariableBinding::fixForHole(ConstraintSystem &cs) const {
     // If the whole body is being ignored due to a pre-check failure,
     // let's not record a fix about result type since there is
     // just not enough context to infer it without a body.
-    if (cs.hasFixFor(cs.getConstraintLocator(closure->getBody()),
+    if (cs.hasFixFor(cs.getConstraintLocator(closure),
                      FixKind::IgnoreInvalidResultBuilderBody))
       return None;
 

lib/Sema/CSDiagnostics.cpp

@@ -7070,3 +7070,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

@@ -2311,6 +2311,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

@@ -1580,7 +1580,7 @@ bool IgnoreInvalidResultBuilderBody::diagnose(const Solution &solution,
     return true; // Already diagnosed by `matchResultBuilder`.
   }
 
-  auto *S = getAnchor().get<Stmt *>();
+  auto *S = castToExpr<ClosureExpr>(getAnchor())->getBody();
 
   class PreCheckWalker : public ASTWalker {
     DeclContext *DC;
@@ -1645,3 +1645,16 @@ AllowRefToInvalidDecl::create(ConstraintSystem &cs,
                               ConstraintLocator *locator) {
   return new (cs.getAllocator()) AllowRefToInvalidDecl(cs, locator);
 }
+
+bool IgnoreResultBuilderWithReturnStmts::diagnose(const Solution &solution,
+                                                  bool asNote) const {
+  InvalidReturnInResultBuilderBody failure(solution, BuilderType, getLocator());
+  return failure.diagnose(asNote);
+}
+
+IgnoreResultBuilderWithReturnStmts *
+IgnoreResultBuilderWithReturnStmts::create(ConstraintSystem &cs, Type builderTy,
+                                           ConstraintLocator *locator) {
+  return new (cs.getAllocator())
+      IgnoreResultBuilderWithReturnStmts(cs, builderTy, locator);
+}

test/Constraints/result_builder_diags.swift

@@ -286,6 +286,21 @@ struct MyTuplifiedStruct {
   }
 }
 
+func test_invalid_return_type_in_body() {
+  tuplify(true) { _ -> (Void, Int) in
+    tuplify(false) { condition in
+      if condition {
+        return 42 // expected-error {{cannot use explicit 'return' statement in the body of result builder 'TupleBuilder'}}
+        // expected-note@-1 {{remove 'return' statements to apply the result builder}} {{9-16=}}
+      } else {
+        1
+      }
+    }
+
+    42
+  }
+}
+
 // Check that we're performing syntactic use diagnostics.
 func acceptMetatype<T>(_: T.Type) -> Bool { true }