Merge pull request #33509 from xedin/rdar-65320500

[FunctionBuilders] Implement graceful handling of pre-check failures

Author: xedin

include/swift/AST/DiagnosticEngine.h

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

@@ -1843,6 +1843,7 @@ class ValueWitnessRequest
 
 struct PreCheckFunctionBuilderDescriptor {
   AnyFunctionRef Fn;
+  bool SuppressDiagnostics;
 
 private:
   // NOTE: Since source tooling (e.g. code completion) might replace the body,
@@ -1852,8 +1853,8 @@ struct PreCheckFunctionBuilderDescriptor {
   BraceStmt *Body;
 
 public:
-  PreCheckFunctionBuilderDescriptor(AnyFunctionRef Fn)
-      : Fn(Fn), Body(Fn.getBody()) {}
+  PreCheckFunctionBuilderDescriptor(AnyFunctionRef Fn, bool suppressDiagnostics)
+      : Fn(Fn), SuppressDiagnostics(suppressDiagnostics), Body(Fn.getBody()) {}
 
   friend llvm::hash_code
   hash_value(const PreCheckFunctionBuilderDescriptor &owner) {

lib/Sema/BuilderTransform.cpp

@@ -1489,9 +1489,11 @@ 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{AnyFunctionRef(func)};
-  switch (evaluateOrDefault(
-              ctx.evaluator, request, FunctionBuilderBodyPreCheck::Error)) {
+  auto request =
+      PreCheckFunctionBuilderRequest{{AnyFunctionRef(func),
+                                      /*SuppressDiagnostics=*/false}};
+  switch (evaluateOrDefault(ctx.evaluator, request,
+                            FunctionBuilderBodyPreCheck::Error)) {
   case FunctionBuilderBodyPreCheck::Okay:
     // If the pre-check was okay, apply the function-builder transform.
     break;
@@ -1623,16 +1625,24 @@ ConstraintSystem::matchFunctionBuilder(
 
   // Pre-check the body: pre-check any expressions in it and look
   // for return statements.
-  auto request = PreCheckFunctionBuilderRequest{fn};
+  auto request =
+      PreCheckFunctionBuilderRequest{{fn, /*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.
-    return getTypeMatchFailure(locator);
+  case FunctionBuilderBodyPreCheck::Error: {
+    if (!shouldAttemptFixes())
+      return getTypeMatchFailure(locator);
+
+    if (recordFix(IgnoreInvalidFunctionBuilderBody::create(
+            *this, getConstraintLocator(fn.getBody()))))
+      return getTypeMatchFailure(locator);
+
+    return getTypeMatchSuccess();
+  }
 
   case FunctionBuilderBodyPreCheck::HasReturnStmt:
     // If the body has a return statement, suppress the transform but
@@ -1709,14 +1719,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; }
 
@@ -1740,16 +1753,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 {
@@ -1782,11 +1807,15 @@ FunctionBuilderBodyPreCheck PreCheckFunctionBuilderRequest::evaluate(
           owner.Fn.getAbstractClosureExpr()))
     skipPrecheck = shouldTypeCheckInEnclosingExpression(closure);
 
-  return PreCheckFunctionBuilderApplication(owner.Fn, false).run();
+  return PreCheckFunctionBuilderApplication(
+             owner.Fn, /*skipPrecheck=*/false,
+             /*suppressDiagnostics=*/owner.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

@@ -1206,7 +1206,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 = castToExpr<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->getBody()),
+                          FixKind::IgnoreInvalidFunctionBuilderBody))
+          fix = SpecifyClosureReturnType::create(cs, dstLocator);
       } else if (srcLocator->directlyAt<ObjectLiteralExpr>()) {
         fix = SpecifyObjectLiteralTypeImport::create(cs, dstLocator);
       } else if (srcLocator->isKeyPathRoot()) {

lib/Sema/CSFix.cpp

@@ -1551,3 +1551,47 @@ AllowKeyPathWithoutComponents::create(ConstraintSystem &cs,
                                       ConstraintLocator *locator) {
   return new (cs.getAllocator()) AllowKeyPathWithoutComponents(cs, locator);
 }
+
+bool IgnoreInvalidFunctionBuilderBody::diagnose(const Solution &solution,
+                                                bool asNote) const {
+  auto *S = getAnchor().get<Stmt *>();
+
+  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();
+    }
+  };
+
+  PreCheckWalker walker(solution.getDC());
+  S->walk(walker);
+
+  return walker.diagnosed();
+}
+
+IgnoreInvalidFunctionBuilderBody *
+IgnoreInvalidFunctionBuilderBody::create(ConstraintSystem &cs,
+                                         ConstraintLocator *locator) {
+  return new (cs.getAllocator()) IgnoreInvalidFunctionBuilderBody(cs, locator);
+}

lib/Sema/CSFix.h

@@ -276,6 +276,9 @@ enum class FixKind : uint8_t {
 
   /// Allow key path expressions with no components.
   AllowKeyPathWithoutComponents,
+
+  /// Ignore function builder body which fails `pre-check` call.
+  IgnoreInvalidFunctionBuilderBody,
 };
 
 class ConstraintFix {
@@ -1980,6 +1983,26 @@ class AllowKeyPathWithoutComponents 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

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

lib/Sema/CSSolver.cpp

@@ -1411,7 +1411,8 @@ void ConstraintSystem::solveForCodeCompletion(
     llvm::function_ref<void(const Solution &)> callback) {
   // First, pre-check the expression, validating any types that occur in the
   // expression and folding sequence expressions.
-  if (ConstraintSystem::preCheckExpression(expr, DC))
+  if (ConstraintSystem::preCheckExpression(
+          expr, DC, /*replaceInvalidRefsWithErrors=*/true))
     return;
 
   ConstraintSystemOptions options;

lib/Sema/ConstraintSystem.h

@@ -4955,8 +4955,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

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -553,10 +553,10 @@ TypeChecker::getTypeOfCompletionOperator(DeclContext *DC, Expr *LHS,
   // We allocate these expressions on the stack because we know they can't
   // escape and there isn't a better way to allocate scratch Expr nodes.
   UnresolvedDeclRefExpr UDRE(DeclNameRef(opName), refKind, DeclNameLoc(Loc));
-  auto *opExpr = TypeChecker::resolveDeclRefExpr(&UDRE, DC);
+  auto *opExpr = TypeChecker::resolveDeclRefExpr(
+      &UDRE, DC, /*replaceInvalidRefsWithErrors=*/true);
 
   switch (refKind) {
-
   case DeclRefKind::PostfixOperator: {
     // (postfix_unary_expr
     //   (declref_expr name=<opName>)
@@ -611,7 +611,8 @@ static Optional<Type> getTypeOfCompletionContextExpr(
                         CompletionTypeCheckKind kind,
                         Expr *&parsedExpr,
                         ConcreteDeclRef &referencedDecl) {
-  if (constraints::ConstraintSystem::preCheckExpression(parsedExpr, DC))
+  if (constraints::ConstraintSystem::preCheckExpression(
+          parsedExpr, DC, /*replaceInvalidRefsWithErrors=*/true))
     return None;
 
   switch (kind) {