[FunctionBuilders] Implement graceful handling of pre-check failures

Function builder bodies are applied in several stages. First of them is `pre-check`
stage which is, among other things, responsible for member lookup and reference
resolution, if that action fails it would be diagnosed inline and affected
expression is going to be replaced with `ErrorExpr`.

This works fine in context of regular expressions but not function builders, because
matching of function builder bodies happens in the middle of type-checking of
enclosing context, so solver needs to implement graceful handling of such problems
and delay diagnostics until solver reaches a solution.

To accommodate that `PreCheckFunctionBuilderRequest`,
`ConstraintSystem::preCheckExpression`, and `TypeChecker::resolveDeclRefExpr`
have been adjusted to know explicitly whether diagnosing and AST mutation is permitted.
Solver is adjusted to ignore closure bodies which fail pre-check in their entirety.

Resolves: rdar://problem/65320500

Author: xedin

include/swift/AST/DiagnosticEngine.h

@@ -1025,6 +1025,12 @@ namespace swift {
       }
     }
 
+    bool hasDiagnostics() const {
+      return std::distance(Engine.TentativeDiagnostics.begin() +
+                               PrevDiagnostics,
+                           Engine.TentativeDiagnostics.end()) > 0;
+    }
+
     /// Abort and close this transaction and erase all diagnostics
     /// record while it was open.
     void abort() {

include/swift/AST/TypeCheckRequests.h

@@ -1822,7 +1822,8 @@ enum class FunctionBuilderBodyPreCheck : uint8_t {
 class PreCheckFunctionBuilderRequest
     : public SimpleRequest<PreCheckFunctionBuilderRequest,
                            FunctionBuilderBodyPreCheck(AnyFunctionRef,
-                                                       BraceStmt *),
+                                                       BraceStmt *,
+                                                       bool),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -1832,7 +1833,8 @@ class PreCheckFunctionBuilderRequest
 
   // Evaluation.
   FunctionBuilderBodyPreCheck evaluate(Evaluator &evaluator, AnyFunctionRef fn,
-                                       BraceStmt *body) const;
+                                       BraceStmt *body,
+                                       bool suppressDiagnostics) const;
 
 public:
   // Separate caching.

lib/Sema/BuilderTransform.cpp

@@ -1374,9 +1374,10 @@ Optional<BraceStmt *> TypeChecker::applyFunctionBuilderBodyTransform(
   // If we encountered an error or there was an explicit result type,
   // bail out and report that to the caller.
   auto &ctx = func->getASTContext();
-  auto request = PreCheckFunctionBuilderRequest{func, func->getBody()};
-  switch (evaluateOrDefault(
-              ctx.evaluator, request, FunctionBuilderBodyPreCheck::Error)) {
+  auto request = PreCheckFunctionBuilderRequest{func, func->getBody(),
+                                                /*suppressDiagnostics=*/false};
+  switch (evaluateOrDefault(ctx.evaluator, request,
+                            FunctionBuilderBodyPreCheck::Error)) {
   case FunctionBuilderBodyPreCheck::Okay:
     // If the pre-check was okay, apply the function-builder transform.
     break;
@@ -1524,16 +1525,27 @@ ConstraintSystem::matchFunctionBuilder(
 
   // Pre-check the body: pre-check any expressions in it and look
   // for return statements.
-  auto request = PreCheckFunctionBuilderRequest{fn, fn.getBody()};
+  auto request = PreCheckFunctionBuilderRequest{fn, fn.getBody(),
+                                                /*SuppressDiagnostics=*/true};
   switch (evaluateOrDefault(getASTContext().evaluator, request,
                             FunctionBuilderBodyPreCheck::Error)) {
   case FunctionBuilderBodyPreCheck::Okay:
     // If the pre-check was okay, apply the function-builder transform.
     break;
 
-  case FunctionBuilderBodyPreCheck::Error:
-    // If the pre-check had an error, flag that.
+  case FunctionBuilderBodyPreCheck::Error: {
+    if (!shouldAttemptFixes())
+      return getTypeMatchFailure(locator);
+
+    if (auto *closure =
+            dyn_cast_or_null<ClosureExpr>(fn.getAbstractClosureExpr())) {
+      auto failed = recordFix(IgnoreInvalidFunctionBuilderBody::create(
+          *this, getConstraintLocator(closure)));
+      return failed ? getTypeMatchFailure(locator) : getTypeMatchSuccess();
+    }
+
     return getTypeMatchFailure(locator);
+  }
 
   case FunctionBuilderBodyPreCheck::HasReturnStmt:
     // If the body has a return statement, suppress the transform but
@@ -1626,14 +1638,17 @@ namespace {
 class PreCheckFunctionBuilderApplication : public ASTWalker {
   AnyFunctionRef Fn;
   bool SkipPrecheck = false;
+  bool SuppressDiagnostics = false;
   std::vector<ReturnStmt *> ReturnStmts;
   bool HasError = false;
 
   bool hasReturnStmt() const { return !ReturnStmts.empty(); }
 
 public:
-  PreCheckFunctionBuilderApplication(AnyFunctionRef fn, bool skipPrecheck)
-    : Fn(fn), SkipPrecheck(skipPrecheck) {}
+  PreCheckFunctionBuilderApplication(AnyFunctionRef fn, bool skipPrecheck,
+                                     bool suppressDiagnostics)
+      : Fn(fn), SkipPrecheck(skipPrecheck),
+        SuppressDiagnostics(suppressDiagnostics) {}
 
   const std::vector<ReturnStmt *> getReturnStmts() const { return ReturnStmts; }
 
@@ -1657,16 +1672,28 @@ class PreCheckFunctionBuilderApplication : public ASTWalker {
   }
 
   std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+    if (SkipPrecheck)
+      return std::make_pair(false, E);
+
     // Pre-check the expression.  If this fails, abort the walk immediately.
     // Otherwise, replace the expression with the result of pre-checking.
     // In either case, don't recurse into the expression.
-    if (!SkipPrecheck &&
-        ConstraintSystem::preCheckExpression(E, /*DC*/ Fn.getAsDeclContext())) {
-      HasError = true;
-      return std::make_pair(false, nullptr);
-    }
+    {
+      auto *DC = Fn.getAsDeclContext();
+      auto &diagEngine = DC->getASTContext().Diags;
+
+      // Suppress any diangostics which could be produced by this expression.
+      DiagnosticTransaction transaction(diagEngine);
 
-    return std::make_pair(false, E);
+      HasError |= ConstraintSystem::preCheckExpression(
+          E, DC, /*replaceInvalidRefsWithErrors=*/false);
+      HasError |= transaction.hasDiagnostics();
+
+      if (SuppressDiagnostics)
+        transaction.abort();
+
+      return std::make_pair(false, HasError ? nullptr : E);
+    }
   }
 
   std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
@@ -1692,18 +1719,20 @@ class PreCheckFunctionBuilderApplication : public ASTWalker {
 
 FunctionBuilderBodyPreCheck
 PreCheckFunctionBuilderRequest::evaluate(Evaluator &eval, AnyFunctionRef fn,
-                                         BraceStmt *body) const {
+                                         BraceStmt *body,
+                                         bool suppressDiagnostics) const {
   // NOTE: 'body' is passed only for the request evaluater caching key.
   // Since source tooling (e.g. code completion) might replace the body,
   // the function alone is not sufficient for the key.
   assert(fn.getBody() == body &&
          "body must be the current body of the function");
 
-  return PreCheckFunctionBuilderApplication(fn, false).run();
+  return PreCheckFunctionBuilderApplication(fn, false, suppressDiagnostics).run();
 }
 
 std::vector<ReturnStmt *> TypeChecker::findReturnStatements(AnyFunctionRef fn) {
-  PreCheckFunctionBuilderApplication precheck(fn, true);
+  PreCheckFunctionBuilderApplication precheck(fn, /*skipPreCheck=*/true,
+                                              /*SuppressDiagnostics=*/true);
   (void)precheck.run();
   return precheck.getReturnStmts();
 }

lib/Sema/CSBindings.cpp

@@ -1106,7 +1106,14 @@ bool TypeVariableBinding::attempt(ConstraintSystem &cs) const {
       } else if (TypeVar->getImpl().isClosureParameterType()) {
         fix = SpecifyClosureParameterType::create(cs, dstLocator);
       } else if (TypeVar->getImpl().isClosureResultType()) {
-        fix = SpecifyClosureReturnType::create(cs, dstLocator);
+        auto *locator = TypeVar->getImpl().getLocator();
+        auto *closure = cast<ClosureExpr>(locator->getAnchor());
+        // 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),
+                          FixKind::IgnoreInvalidFunctionBuilderBody))
+          fix = SpecifyClosureReturnType::create(cs, dstLocator);
       } else if (srcLocator->getAnchor() &&
                  isa<ObjectLiteralExpr>(srcLocator->getAnchor())) {
         fix = SpecifyObjectLiteralTypeImport::create(cs, dstLocator);

lib/Sema/CSFix.cpp

@@ -1363,3 +1363,48 @@ bool SpecifyKeyPathRootType::diagnose(const Solution &solution,
 
   return failure.diagnose(asNote);
 }
+
+bool IgnoreInvalidFunctionBuilderBody::diagnose(const Solution &solution,
+                                                bool asNote) const {
+  auto *S = cast<ClosureExpr>(getAnchor())->getBody();
+
+  class PreCheckWalker : public ASTWalker {
+    DeclContext *DC;
+    DiagnosticTransaction Transaction;
+
+  public:
+    PreCheckWalker(DeclContext *dc)
+        : DC(dc), Transaction(dc->getASTContext().Diags) {}
+
+    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+      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 function builder pre-check does so as well.
+    std::pair<bool, Pattern *> walkToPatternPre(Pattern *P) override {
+      return std::make_pair(false, P);
+    }
+
+    bool diagnosed() const {
+      return Transaction.hasDiagnostics();
+    }
+  };
+
+  auto &cs = getConstraintSystem();
+  PreCheckWalker walker(cs.DC);
+  S->walk(walker);
+
+  return walker.diagnosed();
+}
+
+IgnoreInvalidFunctionBuilderBody *
+IgnoreInvalidFunctionBuilderBody::create(ConstraintSystem &cs,
+                                         ConstraintLocator *locator) {
+  return new (cs.getAllocator()) IgnoreInvalidFunctionBuilderBody(cs, locator);
+}

lib/Sema/CSFix.h

@@ -264,6 +264,8 @@ enum class FixKind : uint8_t {
   /// Specify key path root type when it cannot be infered from context.
   SpecifyKeyPathRootType,
 
+  /// Ignore function builder body which fails `pre-check` call.
+  IgnoreInvalidFunctionBuilderBody,
 };
 
 class ConstraintFix {
@@ -1841,6 +1843,26 @@ class SpecifyKeyPathRootType final : public ConstraintFix {
                                           ConstraintLocator *locator);
 };
 
+class IgnoreInvalidFunctionBuilderBody final : public ConstraintFix {
+  IgnoreInvalidFunctionBuilderBody(ConstraintSystem &cs,
+                                   ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::IgnoreInvalidFunctionBuilderBody, locator) {}
+
+public:
+  std::string getName() const override {
+    return "ignore invalid function builder body";
+  }
+
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  bool diagnoseForAmbiguity(CommonFixesArray commonFixes) const override {
+    return diagnose(*commonFixes.front().first);
+  }
+
+  static IgnoreInvalidFunctionBuilderBody *create(ConstraintSystem &cs,
+                                                  ConstraintLocator *locator);
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSSimplify.cpp

@@ -9712,7 +9712,8 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::SpecifyObjectLiteralTypeImport:
   case FixKind::AllowKeyPathRootTypeMismatch:
   case FixKind::AllowCoercionToForceCast:
-  case FixKind::SpecifyKeyPathRootType: {
+  case FixKind::SpecifyKeyPathRootType:
+  case FixKind::IgnoreInvalidFunctionBuilderBody: {
     return recordFix(fix) ? SolutionKind::Error : SolutionKind::Solved;
   }
 

lib/Sema/ConstraintSystem.h

@@ -4606,8 +4606,12 @@ class ConstraintSystem {
 public:
   /// Pre-check the expression, validating any types that occur in the
   /// expression and folding sequence expressions.
-  static bool preCheckExpression(Expr *&expr, DeclContext *dc);
-        
+  ///
+  /// \param replaceInvalidRefsWithErrors Indicates whether it's allowed
+  /// to replace any discovered invalid member references with `ErrorExpr`.
+  static bool preCheckExpression(Expr *&expr, DeclContext *dc,
+                                 bool replaceInvalidRefsWithErrors);
+
   /// Solve the system of constraints generated from provided target.
   ///
   /// \param target The target that we'll generate constraints from, which