[ConstraintSystem] Use special locator for expr representing return of single expr function

Author: xedin

Diff for: lib/Sema/CSApply.cpp

@@ -7696,10 +7696,18 @@ Expr *ConstraintSystem::applySolution(Solution &solution, Expr *expr,
     if (hadError)
       return nullptr;
   }
-  
+
+  // We are supposed to use contextual type only if it is present and
+  // this expression doesn't represent the implicit return of the single
+  // expression function which got deduced to be `Never`.
+  auto shouldCoerceToContextualType = [&]() {
+    return convertType && !(getType(result)->isUninhabited() &&
+                            getContextualTypePurpose() == CTP_ReturnSingleExpr);
+  };
+
   // If we're supposed to convert the expression to some particular type,
   // do so now.
-  if (convertType) {
+  if (shouldCoerceToContextualType()) {
     result = rewriter.coerceToType(result, convertType,
                                    getConstraintLocator(expr));
     if (!result)

Diff for: lib/Sema/CSSimplify.cpp

@@ -2793,9 +2793,11 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
   }
 
   // Allow '() -> T' to '() -> ()' and '() -> Never' to '() -> T' for closure
-  // literals.
+  // literals and expressions representing an implicit return type of the single
+  // expression functions.
   if (auto elt = locator.last()) {
-    if (elt->getKind() == ConstraintLocator::ClosureResult) {
+    if (elt->getKind() == ConstraintLocator::ClosureResult ||
+        elt->getKind() == ConstraintLocator::SingleExprFuncResultType) {
       if (kind >= ConstraintKind::Subtype &&
           (type1->isUninhabited() || type2->isVoid())) {
         increaseScore(SK_FunctionConversion);

Diff for: lib/Sema/CSSolver.cpp

@@ -1162,16 +1162,16 @@ ConstraintSystem::solveImpl(Expr *&expr,
 
     if (getContextualTypePurpose() == CTP_CallArgument)
       constraintKind = ConstraintKind::ArgumentConversion;
-    else if (getContextualTypePurpose() == CTP_ReturnSingleExpr)
-      constraintKind = ConstraintKind::SingleExpressionFunctionReturnConversion;
 
     // In a by-reference yield, we expect the contextual type to be an
     // l-value type, so the result must be bound to that.
     if (getContextualTypePurpose() == CTP_YieldByReference)
       constraintKind = ConstraintKind::Bind;
 
     auto *convertTypeLocator = getConstraintLocator(
-        getConstraintLocator(expr), ConstraintLocator::ContextualType);
+        expr, getContextualTypePurpose() == CTP_ReturnSingleExpr
+                  ? ConstraintLocator::SingleExprFuncResultType
+                  : ConstraintLocator::ContextualType);
 
     if (allowFreeTypeVariables == FreeTypeVariableBinding::UnresolvedType) {
       convertType = convertType.transform([&](Type type) -> Type {

Diff for: lib/Sema/TypeCheckStmt.cpp

@@ -493,7 +493,8 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
     }
 
     ContextualTypePurpose ctp = CTP_ReturnStmt;
-    if (auto func = dyn_cast_or_null<FuncDecl>(TheFunc->getAbstractFunctionDecl())) {
+    if (auto func =
+            dyn_cast_or_null<FuncDecl>(TheFunc->getAbstractFunctionDecl())) {
       if (func->hasSingleExpressionBody()) {
         ctp = CTP_ReturnSingleExpr;
       }
@@ -502,7 +503,6 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
     auto exprTy = TC.typeCheckExpression(E, DC, TypeLoc::withoutLoc(ResultTy),
                                          ctp,
                                          options);
-    
     RS->setResult(E);
 
     if (!exprTy) {
@@ -1921,7 +1921,7 @@ bool TypeChecker::typeCheckFunctionBodyUntil(FuncDecl *FD,
       // The function returns void.  We don't need an explicit return, no matter
       // what the type of the expression is.  Take the inserted return back out.
       BS->setElement(0, E);
-      // Fall through to type-checking the body as if we were not a single 
+      // Fall through to type-checking the body as if we were not a single
       // expression function.
     }
   }
@@ -1930,6 +1930,22 @@ bool TypeChecker::typeCheckFunctionBodyUntil(FuncDecl *FD,
   SC.EndTypeCheckLoc = EndTypeCheckLoc;
   bool HadError = SC.typeCheckBody(BS);
 
+  // If this was a function with a single expression body, let's see
+  // if implicit return statement came out to be `Never` which means
+  // that we have eagerly converted something like `{ fatalError() }`
+  // into `{ return fatalError() }` that has to be corrected here.
+  if (FD->hasSingleExpressionBody()) {
+    if (auto *stmt = BS->getElement(0).dyn_cast<Stmt *>()) {
+      if (auto *RS = dyn_cast<ReturnStmt>(stmt)) {
+        if (RS->isImplicit() && RS->hasResult()) {
+          auto returnType = RS->getResult()->getType();
+          if (returnType && returnType->isUninhabited())
+            BS->setElement(0, RS->getResult());
+        }
+      }
+    }
+  }
+
   FD->setBody(BS);
   return HadError;
 }