[ResultBuilders] Diagnose pre-check errors inline (#37353)

Not all of the pre-check errors could be diagnosed by re-running
`PreCheckExpression` e.g. key path expressions are mutated to
a particular form after an error has been produced.

To make the behavior consistent, let's allow pre-check to emit
diagnostics and unify pre-check and constraint generation fixes.

Resolves: rdar://77466241
(cherry picked from commit 4b0d9a2)

Author: xedin

include/swift/Sema/CSFix.h

@@ -2183,17 +2183,9 @@ class AllowKeyPathWithoutComponents final : public ConstraintFix {
 
 class IgnoreInvalidResultBuilderBody : public ConstraintFix {
 protected:
-  enum class ErrorInPhase {
-    PreCheck,
-    ConstraintGeneration,
-  };
-
-  ErrorInPhase Phase;
-
-  IgnoreInvalidResultBuilderBody(ConstraintSystem &cs, ErrorInPhase phase,
-                                   ConstraintLocator *locator)
-      : ConstraintFix(cs, FixKind::IgnoreInvalidResultBuilderBody, locator),
-        Phase(phase) {}
+  IgnoreInvalidResultBuilderBody(ConstraintSystem &cs,
+                                 ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::IgnoreInvalidResultBuilderBody, locator) {}
 
 public:
   std::string getName() const override {
@@ -2206,19 +2198,8 @@ class IgnoreInvalidResultBuilderBody : public ConstraintFix {
     return diagnose(*commonFixes.front().first);
   }
 
-  static IgnoreInvalidResultBuilderBody *
-  duringPreCheck(ConstraintSystem &cs, ConstraintLocator *locator) {
-    return create(cs, ErrorInPhase::PreCheck, locator);
-  }
-
-  static IgnoreInvalidResultBuilderBody *
-  duringConstraintGeneration(ConstraintSystem &cs, ConstraintLocator *locator) {
-    return create(cs, ErrorInPhase::ConstraintGeneration, locator);
-  }
-
-private:
-  static IgnoreInvalidResultBuilderBody *
-  create(ConstraintSystem &cs, ErrorInPhase phase, ConstraintLocator *locator);
+  static IgnoreInvalidResultBuilderBody *create(ConstraintSystem &cs,
+                                                ConstraintLocator *locator);
 };
 
 class IgnoreResultBuilderWithReturnStmts final
@@ -2227,8 +2208,7 @@ class IgnoreResultBuilderWithReturnStmts final
 
   IgnoreResultBuilderWithReturnStmts(ConstraintSystem &cs, Type builderTy,
                                      ConstraintLocator *locator)
-      : IgnoreInvalidResultBuilderBody(cs, ErrorInPhase::PreCheck, locator),
-        BuilderType(builderTy) {}
+      : IgnoreInvalidResultBuilderBody(cs, locator), BuilderType(builderTy) {}
 
 public:
   bool diagnose(const Solution &solution, bool asNote = false) const override;

lib/Sema/BuilderTransform.cpp

@@ -1691,27 +1691,28 @@ ConstraintSystem::matchResultBuilder(
   assert(builder->getAttrs().hasAttribute<ResultBuilderAttr>());
 
   if (InvalidResultBuilderBodies.count(fn)) {
-    (void)recordFix(
-        IgnoreInvalidResultBuilderBody::duringConstraintGeneration(
-            *this, getConstraintLocator(fn.getAbstractClosureExpr())));
+    (void)recordFix(IgnoreInvalidResultBuilderBody::create(
+        *this, getConstraintLocator(fn.getAbstractClosureExpr())));
     return getTypeMatchSuccess();
   }
 
   // Pre-check the body: pre-check any expressions in it and look
   // for return statements.
   auto request =
-      PreCheckResultBuilderRequest{{fn, /*SuppressDiagnostics=*/true}};
+      PreCheckResultBuilderRequest{{fn, /*SuppressDiagnostics=*/false}};
   switch (evaluateOrDefault(getASTContext().evaluator, request,
                             ResultBuilderBodyPreCheck::Error)) {
   case ResultBuilderBodyPreCheck::Okay:
     // If the pre-check was okay, apply the result-builder transform.
     break;
 
   case ResultBuilderBodyPreCheck::Error: {
+    InvalidResultBuilderBodies.insert(fn);
+
     if (!shouldAttemptFixes())
       return getTypeMatchFailure(locator);
 
-    if (recordFix(IgnoreInvalidResultBuilderBody::duringPreCheck(
+    if (recordFix(IgnoreInvalidResultBuilderBody::create(
             *this, getConstraintLocator(fn.getAbstractClosureExpr()))))
       return getTypeMatchFailure(locator);
 
@@ -1776,9 +1777,8 @@ ConstraintSystem::matchResultBuilder(
     if (transaction.hasErrors()) {
       InvalidResultBuilderBodies.insert(fn);
 
-      if (recordFix(
-              IgnoreInvalidResultBuilderBody::duringConstraintGeneration(
-                  *this, getConstraintLocator(fn.getAbstractClosureExpr()))))
+      if (recordFix(IgnoreInvalidResultBuilderBody::create(
+              *this, getConstraintLocator(fn.getAbstractClosureExpr()))))
         return getTypeMatchFailure(locator);
 
       return getTypeMatchSuccess();
@@ -1869,7 +1869,7 @@ class PreCheckResultBuilderApplication : public ASTWalker {
       DiagnosticTransaction transaction(diagEngine);
 
       HasError |= ConstraintSystem::preCheckExpression(
-          E, DC, /*replaceInvalidRefsWithErrors=*/false);
+          E, DC, /*replaceInvalidRefsWithErrors=*/true);
       HasError |= transaction.hasErrors();
 
       if (!HasError) {

lib/Sema/CSFix.cpp

@@ -1722,78 +1722,14 @@ AllowKeyPathWithoutComponents::create(ConstraintSystem &cs,
 }
 
 bool IgnoreInvalidResultBuilderBody::diagnose(const Solution &solution,
-                                                bool asNote) const {
-  switch (Phase) {
-  // Handled below
-  case ErrorInPhase::PreCheck:
-    break;
-  case ErrorInPhase::ConstraintGeneration:
-    return true; // Already diagnosed by `matchResultBuilder`.
-  }
-
-  auto *S = castToExpr<ClosureExpr>(getAnchor())->getBody();
-
-  class PreCheckWalker : public ASTWalker {
-    DeclContext *DC;
-    DiagnosticTransaction Transaction;
-    // Check whether expression(s) in the body of the
-    // result builder had any `ErrorExpr`s before pre-check.
-    bool FoundErrorExpr = false;
-
-  public:
-    PreCheckWalker(DeclContext *dc)
-        : DC(dc), Transaction(dc->getASTContext().Diags) {}
-
-    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-      // This has to be checked before `preCheckExpression`
-      // because pre-check would convert invalid references
-      // into `ErrorExpr`s.
-      E->forEachChildExpr([&](Expr *expr) {
-        FoundErrorExpr |= isa<ErrorExpr>(expr);
-        return FoundErrorExpr ? nullptr : expr;
-      });
-
-      auto hasError = ConstraintSystem::preCheckExpression(
-          E, DC, /*replaceInvalidRefsWithErrors=*/true);
-
-      return std::make_pair(false, hasError ? nullptr : E);
-    }
-
-    std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
-      return std::make_pair(true, S);
-    }
-
-    // Ignore patterns because result builder pre-check does so as well.
-    std::pair<bool, Pattern *> walkToPatternPre(Pattern *P) override {
-      return std::make_pair(false, P);
-    }
-
-    bool diagnosed() const {
-      // pre-check produced an error.
-      if (Transaction.hasErrors())
-        return true;
-
-      // If there were `ErrorExpr`s before pre-check
-      // they should have been diagnosed already by parser.
-      if (FoundErrorExpr) {
-        auto &DE = DC->getASTContext().Diags;
-        return DE.hadAnyError();
-      }
-
-      return false;
-    }
-  };
-
-  PreCheckWalker walker(solution.getDC());
-  S->walk(walker);
-
-  return walker.diagnosed();
+                                              bool asNote) const {
+  return true; // Already diagnosed by `matchResultBuilder`.
 }
 
-IgnoreInvalidResultBuilderBody *IgnoreInvalidResultBuilderBody::create(
-    ConstraintSystem &cs, ErrorInPhase phase, ConstraintLocator *locator) {
-  return new (cs.getAllocator())
-      IgnoreInvalidResultBuilderBody(cs, phase, locator);
+IgnoreInvalidResultBuilderBody *
+IgnoreInvalidResultBuilderBody::create(ConstraintSystem &cs,
+                                       ConstraintLocator *locator) {
+  return new (cs.getAllocator()) IgnoreInvalidResultBuilderBody(cs, locator);
 }
 
 bool SpecifyContextualTypeForNil::diagnose(const Solution &solution,

test/Constraints/rdar65320500.swift

@@ -24,20 +24,20 @@ func test(_: Int) -> Bool {
 }
 
 test_builder {
-  let totalSeconds = 42000
-  test(totalseconds / 3600) // expected-error {{cannot find 'totalseconds' in scope}}
+  let totalSeconds = 42000 // expected-note {{'totalSeconds' declared here}}
+  test(totalseconds / 3600) // expected-error {{cannot find 'totalseconds' in scope; did you mean 'totalSeconds'?}}
 }
 
 test_builder {
   test(doesntExist()) // expected-error {{cannot find 'doesntExist' in scope}}
 
-  if let result = doesntExist() { // expected-error {{cannot find 'doesntExist' in scope}}
+  if let result = doesntExist() {
   }
 
-  if bar = test(42) {} // expected-error {{cannot find 'bar' in scope}}
+  if bar = test(42) {}
 
-  let foo = bar() // expected-error {{cannot find 'bar' in scope}}
+  let foo = bar()
 
-  switch (doesntExist()) { // expected-error {{cannot find 'doesntExist' in scope}}
+  switch (doesntExist()) {
   }
 }

test/Constraints/result_builder_diags.swift

@@ -142,6 +142,12 @@ func testOverloading(name: String) {
     }
   }
 
+  _ = overloadedTuplify(true) { cond in
+    if cond {
+      print(\"hello") // expected-error {{invalid component of Swift key path}}
+    }
+  }
+
   let _: A = a1
 
   _ = overloadedTuplify(true) { b in // expected-error {{ambiguous use of 'overloadedTuplify(_:body:)'}}
@@ -461,7 +467,7 @@ struct TestConstraintGenerationErrors {
   func buildTupleClosure() {
     tuplify(true) { _ in
       let a = nothing // expected-error {{cannot find 'nothing' in scope}}
-      String(nothing) // expected-error {{cannot find 'nothing' in scope}}
+      String(nothing)
     }
   }
 }