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

Author: xedin

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/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);