[Diagnostics] Diagnose inability to infer (complex) closure return type

include/swift/AST/DiagnosticsSema.def

@@ -302,11 +302,11 @@ ERROR(cannot_invoke_closure_type,none,
 ERROR(cannot_infer_closure_type,none,
       "unable to infer closure type in the current context", ())
 ERROR(cannot_infer_closure_result_type,none,
-      "unable to infer complex closure return type; "
-      "add explicit type to disambiguate", ())
-FIXIT(insert_closure_return_type,
-      "%select{| () }1-> %0 %select{|in }1",
-      (Type, bool))
+      "unable to infer%select{ complex|}0 closure return type; "
+      "add explicit type to disambiguate", (bool))
+FIXIT(insert_closure_return_type_placeholder,
+      "%select{| () }0-> <#Result#> %select{|in }0",
+      (bool))
 
 ERROR(incorrect_explicit_closure_result,none,
       "declared closure result %0 is incompatible with contextual type %1",

lib/Sema/CSBindings.cpp

@@ -999,6 +999,11 @@ bool TypeVariableBinding::attempt(ConstraintSystem &cs) const {
             cs.getConstraintLocator(dstLocator->getAnchor(), path.drop_back()));
         if (cs.recordFix(fix))
           return true;
+      } else if (TypeVar->getImpl().isClosureResultType()) {
+        auto *fix = SpecifyClosureReturnType::create(
+            cs, TypeVar->getImpl().getLocator());
+        if (cs.recordFix(fix))
+          return true;
       }
     }
   }

lib/Sema/CSDiag.cpp

@@ -220,11 +220,6 @@ class FailureDiagnosis :public ASTVisitor<FailureDiagnosis, /*exprresult*/bool>{
   std::pair<Type, ContextualTypePurpose>
   validateContextualType(Type contextualType, ContextualTypePurpose CTP);
 
-  /// Check the specified closure to see if it is a multi-statement closure with
-  /// an uninferred type.  If so, diagnose the problem with an error and return
-  /// true.
-  bool diagnoseAmbiguousMultiStatementClosure(ClosureExpr *closure);
-
   /// Given a result of name lookup that had no viable results, diagnose the
   /// unviable ones.
   void diagnoseUnviableLookupResults(MemberLookupResult &lookupResults,
@@ -2694,138 +2689,6 @@ FailureDiagnosis::validateContextualType(Type contextualType,
   return {contextualType, CTP};
 }
 
-/// Check the specified closure to see if it is a multi-statement closure with
-/// an uninferred type.  If so, diagnose the problem with an error and return
-/// true.
-bool FailureDiagnosis::
-diagnoseAmbiguousMultiStatementClosure(ClosureExpr *closure) {
-  if (closure->hasSingleExpressionBody() ||
-      closure->hasExplicitResultType())
-    return false;
-
-  auto closureType = CS.getType(closure)->getAs<AnyFunctionType>();
-  if (!closureType ||
-      !(closureType->getResult()->hasUnresolvedType() ||
-        closureType->getResult()->hasTypeVariable()))
-    return false;
-
-  // Okay, we have a multi-statement closure expr that has no inferred result,
-  // type, in the context of a larger expression.  The user probably expected
-  // the compiler to infer the result type of the closure from the body of the
-  // closure, which Swift doesn't do for multi-statement closures.  Try to be
-  // helpful by digging into the body of the closure, looking for a return
-  // statement, and inferring the result type from it.  If we can figure that
-  // out, we can produce a fixit hint.
-  class ReturnStmtFinder : public ASTWalker {
-    SmallVectorImpl<ReturnStmt*> &returnStmts;
-  public:
-    ReturnStmtFinder(SmallVectorImpl<ReturnStmt*> &returnStmts)
-      : returnStmts(returnStmts) {}
-
-    // Walk through statements, so we find returns hiding in if/else blocks etc.
-    std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
-      // Keep track of any return statements we find.
-      if (auto RS = dyn_cast<ReturnStmt>(S))
-        returnStmts.push_back(RS);
-      return { true, S };
-    }
-
-    // Don't walk into anything else, since they cannot contain statements
-    // that can return from the current closure.
-    std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-      return { false, E };
-    }
-    std::pair<bool, Pattern*> walkToPatternPre(Pattern *P) override {
-      return { false, P };
-    }
-    bool walkToDeclPre(Decl *D) override { return false; }
-    bool walkToTypeLocPre(TypeLoc &TL) override { return false; }
-    bool walkToTypeReprPre(TypeRepr *T) override { return false; }
-    bool walkToParameterListPre(ParameterList *PL) override { return false; }
-  };
-
-  SmallVector<ReturnStmt*, 4> Returns;
-  closure->getBody()->walk(ReturnStmtFinder(Returns));
-
-  // If we found a return statement inside of the closure expression, then go
-  // ahead and type check the body to see if we can determine a type.
-  for (auto RS : Returns) {
-    llvm::SaveAndRestore<DeclContext *> SavedDC(CS.DC, closure);
-
-    // Otherwise, we're ok to type check the subexpr.
-    Type resultType;
-    if (RS->hasResult()) {
-      auto resultExpr = RS->getResult();
-      ConcreteDeclRef decl = nullptr;
-
-      // If return expression uses closure parameters, which have/are
-      // type variables, such means that we won't be able to
-      // type-check result correctly and, unfortunately,
-      // we are going to leak type variables from the parent
-      // constraint system through declaration types.
-      bool hasUnresolvedParams = false;
-      resultExpr->forEachChildExpr([&](Expr *childExpr) -> Expr *{
-        if (auto DRE = dyn_cast<DeclRefExpr>(childExpr)) {
-          if (auto param = dyn_cast<ParamDecl>(DRE->getDecl())) {
-            auto paramType =
-                param->hasInterfaceType() ? param->getType() : Type();
-            if (!paramType || paramType->hasTypeVariable()) {
-              hasUnresolvedParams = true;
-              return nullptr;
-            }
-          }
-        }
-        return childExpr;
-      });
-
-      if (hasUnresolvedParams)
-        continue;
-
-      ConstraintSystem::preCheckExpression(resultExpr, CS.DC, &CS);
-
-      // Obtain type of the result expression without applying solutions,
-      // because otherwise this might result in leaking of type variables,
-      // since we are not resetting result statement and if expression is
-      // successfully type-checked its type cleanup is going to be disabled
-      // (we are allowing unresolved types), and as a side-effect it might
-      // also be transformed e.g. OverloadedDeclRefExpr -> DeclRefExpr.
-      auto type = TypeChecker::getTypeOfExpressionWithoutApplying(
-          resultExpr, CS.DC, decl, FreeTypeVariableBinding::UnresolvedType);
-      if (type)
-        resultType = type->getRValueType();
-    }
-
-    // If we found a type, presuppose it was the intended result and insert a
-    // fixit hint.
-    if (resultType && !isUnresolvedOrTypeVarType(resultType)) {
-      // If there is a location for an 'in' token, then the argument list was
-      // specified somehow but no return type was.  Insert a "-> ReturnType "
-      // before the in token.
-      if (closure->getInLoc().isValid()) {
-        diagnose(closure->getLoc(), diag::cannot_infer_closure_result_type)
-            .fixItInsert(closure->getInLoc(), diag::insert_closure_return_type,
-                         resultType, /*argListSpecified*/ false);
-        return true;
-      }
-
-      // Otherwise, the closure must take zero arguments.  We know this
-      // because the if one or more argument is specified, a multi-statement
-      // closure *must* name them, or explicitly ignore them with "_ in".
-      //
-      // As such, we insert " () -> ReturnType in " right after the '{' that
-      // starts the closure body.
-      diagnose(closure->getLoc(), diag::cannot_infer_closure_result_type)
-          .fixItInsertAfter(closure->getBody()->getLBraceLoc(),
-                            diag::insert_closure_return_type, resultType,
-                            /*argListSpecified*/ true);
-      return true;
-    }
-  }
-
-  diagnose(closure->getLoc(), diag::cannot_infer_closure_result_type);
-  return true;
-}
-
 /// Emit an ambiguity diagnostic about the specified expression.
 void FailureDiagnosis::diagnoseAmbiguity(Expr *E) {
   if (auto *assignment = dyn_cast<AssignExpr>(E)) {
@@ -2855,11 +2718,6 @@ void FailureDiagnosis::diagnoseAmbiguity(Expr *E) {
   // Unresolved/Anonymous ClosureExprs are common enough that we should give
   // them tailored diagnostics.
   if (auto CE = dyn_cast<ClosureExpr>(E->getValueProvidingExpr())) {
-    // If this is a multi-statement closure with no explicit result type, emit
-    // a note to clue the developer in.
-    if (diagnoseAmbiguousMultiStatementClosure(CE))
-      return;
-
     diagnose(E->getLoc(), diag::cannot_infer_closure_type)
       .highlight(E->getSourceRange());
     return;
@@ -2901,22 +2759,6 @@ void FailureDiagnosis::diagnoseAmbiguity(Expr *E) {
     return;
   }
 
-  // A very common cause of this diagnostic is a situation where a closure expr
-  // has no inferred type, due to being a multiline closure.  Check to see if
-  // this is the case and (if so), speculatively diagnose that as the problem.
-  bool didDiagnose = false;
-  E->forEachChildExpr([&](Expr *subExpr) -> Expr*{
-    auto closure = dyn_cast<ClosureExpr>(subExpr);
-    if (!didDiagnose && closure)
-      didDiagnose = diagnoseAmbiguousMultiStatementClosure(closure);
-
-    return subExpr;
-  });
-
-  if (didDiagnose) return;
-
-
-
   // Attempt to re-type-check the entire expression, allowing ambiguity, but
   // ignoring a contextual type.
   if (expr == E) {

lib/Sema/CSDiagnostics.cpp

@@ -5915,3 +5915,31 @@ bool MissingContextualBaseInMemberRefFailure::diagnoseAsError() {
       .highlight(anchor->getSourceRange());
   return true;
 }
+
+bool UnableToInferClosureReturnType::diagnoseAsError() {
+  auto *closure = cast<ClosureExpr>(getRawAnchor());
+
+  auto diagnostic =
+      emitDiagnostic(closure->getLoc(),
+                     diag::cannot_infer_closure_result_type,
+                     closure->hasSingleExpressionBody());
+
+  // If there is a location for an 'in' token, then the argument list was
+  // specified somehow but no return type was.  Insert a "-> ReturnType "
+  // before the in token.
+  if (closure->getInLoc().isValid()) {
+    diagnostic.fixItInsert(closure->getInLoc(),
+                           diag::insert_closure_return_type_placeholder,
+                           /*argListSpecified=*/false);
+  } else if (closure->getParameters()->size() == 0) {
+    // Otherwise, the closure must take zero arguments.
+    //
+    // As such, we insert " () -> ReturnType in " right after the '{' that
+    // starts the closure body.
+    diagnostic.fixItInsertAfter(closure->getBody()->getLBraceLoc(),
+                                diag::insert_closure_return_type_placeholder,
+                                /*argListSpecified=*/true);
+  }
+
+  return true;
+}

lib/Sema/CSDiagnostics.h

@@ -1912,6 +1912,15 @@ class MissingContextualBaseInMemberRefFailure final : public FailureDiagnostic {
   bool diagnoseAsError();
 };
 
+class UnableToInferClosureReturnType final : public FailureDiagnostic {
+public:
+  UnableToInferClosureReturnType(ConstraintSystem &cs,
+                                 ConstraintLocator *locator)
+      : FailureDiagnostic(cs, locator) {}
+
+  bool diagnoseAsError();
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSFix.cpp

@@ -1155,3 +1155,15 @@ SpecifyBaseTypeForContextualMember *SpecifyBaseTypeForContextualMember::create(
   return new (cs.getAllocator())
       SpecifyBaseTypeForContextualMember(cs, member, locator);
 }
+
+bool SpecifyClosureReturnType::diagnose(bool asNote) const {
+  auto &cs = getConstraintSystem();
+  UnableToInferClosureReturnType failure(cs, getLocator());
+  return failure.diagnose(asNote);
+}
+
+SpecifyClosureReturnType *
+SpecifyClosureReturnType::create(ConstraintSystem &cs,
+                                 ConstraintLocator *locator) {
+  return new (cs.getAllocator()) SpecifyClosureReturnType(cs, locator);
+}

lib/Sema/CSFix.h

@@ -231,6 +231,10 @@ enum class FixKind : uint8_t {
   /// Base type in reference to the contextual member e.g. `.foo` couldn't be
   /// inferred and has to be specified explicitly.
   SpecifyBaseTypeForContextualMember,
+
+  /// Closure return type has to be explicitly specified because it can't be
+  /// inferred in current context e.g. because it's a multi-statement closure.
+  SpecifyClosureReturnType,
 };
 
 class ConstraintFix {
@@ -1605,6 +1609,21 @@ class SpecifyBaseTypeForContextualMember final : public ConstraintFix {
   create(ConstraintSystem &cs, DeclNameRef member, ConstraintLocator *locator);
 };
 
+class SpecifyClosureReturnType final : public ConstraintFix {
+  SpecifyClosureReturnType(ConstraintSystem &cs, ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::SpecifyClosureReturnType, locator) {}
+
+public:
+  std::string getName() const {
+    return "specify closure return type";
+  }
+
+  bool diagnose(bool asNote = false) const;
+
+  static SpecifyClosureReturnType *create(ConstraintSystem &cs,
+                                          ConstraintLocator *locator);
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSGen.cpp

@@ -2430,7 +2430,8 @@ namespace {
           CS.getConstraintLocator(expr, ConstraintLocator::ClosureResult);
 
         if (expr->hasEmptyBody()) {
-          resultTy = CS.createTypeVariable(locator, 0);
+          resultTy = CS.createTypeVariable(
+              locator, expr->hasSingleExpressionBody() ? 0 : TVO_CanBindToHole);
 
           // Closures with empty bodies should be inferred to return
           // ().
@@ -2442,7 +2443,8 @@ namespace {
         } else {
           // If no return type was specified, create a fresh type
           // variable for it.
-          resultTy = CS.createTypeVariable(locator, 0);
+          resultTy = CS.createTypeVariable(
+              locator, expr->hasSingleExpressionBody() ? 0 : TVO_CanBindToHole);
 
           if (closureHasNoResult(expr)) {
             // Allow it to default to () if there are no return statements.

lib/Sema/CSSimplify.cpp

@@ -8596,6 +8596,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::AllowMutatingMemberOnRValueBase:
   case FixKind::AllowTupleSplatForSingleParameter:
   case FixKind::AllowInvalidUseOfTrailingClosure:
+  case FixKind::SpecifyClosureReturnType:
     llvm_unreachable("handled elsewhere");
   }
 

lib/Sema/ConstraintSystem.h

@@ -297,6 +297,9 @@ class TypeVariableType::Implementation {
   /// Retrieve the generic parameter opened by this type variable.
   GenericTypeParamType *getGenericParameter() const;
 
+  /// Determine whether this type variable represents a closure result type.
+  bool isClosureResultType() const;
+
   /// Retrieve the representative of the equivalence class to which this
   /// type variable belongs.
   ///

lib/Sema/TypeCheckConstraints.cpp

@@ -83,6 +83,14 @@ TypeVariableType::Implementation::getGenericParameter() const {
   return locator ? locator->getGenericParameter() : nullptr;
 }
 
+bool TypeVariableType::Implementation::isClosureResultType() const {
+  if (!(locator && locator->getAnchor()))
+    return false;
+
+  return isa<ClosureExpr>(locator->getAnchor()) &&
+         locator->isLastElement<LocatorPathElt::ClosureResult>();
+}
+
 void *operator new(size_t bytes, ConstraintSystem& cs,
                    size_t alignment) {
   return cs.getAllocator().Allocate(bytes, alignment);