[ResultBuilders] Use new fix to detect and diagnose use of return statements in a result builder body

Author: xedin

lib/Sema/BuilderTransform.cpp

@@ -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

@@ -1424,7 +1424,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/CSFix.cpp

@@ -1579,7 +1579,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;

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 {{application of result builder 'TupleBuilder' disabled by explicit 'return' statement}}
+        // 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 }