perform return typechecking instead of assignment for statement function body

5 files changed, 14 insertions, 8 deletions

diff --git a/include/flang/Basic/DiagnosticSemaKinds.td b/include/flang/Basic/DiagnosticSemaKinds.td
index c00fc612dc..1cde03b208 100644
--- a/include/flang/Basic/DiagnosticSemaKinds.td
+++ b/include/flang/Basic/DiagnosticSemaKinds.td
@@ -41,6 +41,8 @@ def err_typecheck_assign_incompatible : Error<
   "assigning to %0 from incompatible type %1">;
 def err_typecheck_initialization_incompatible : Error<
   "initializing %0 with an expression of incompatible type %1">;
+def err_typecheck_return_incompatible : Error<
+  "returning %0 from a function with incompatible result type %1">;
 def err_typecheck_passing_incompatible : Error<
   "passing %0 to parameter of incompatible type %1">;
 def err_typecheck_passing_incompatible_named_arg : Error<
diff --git a/lib/Sema/Sema.cpp b/lib/Sema/Sema.cpp
index 1d4a8e3e8d..757b4ec469 100644
--- a/lib/Sema/Sema.cpp
+++ b/lib/Sema/Sema.cpp
@@ -451,10 +451,11 @@ FunctionDecl *Sema::ActOnStatementFunction(ASTContext &C,
 void Sema::ActOnStatementFunctionBody(SourceLocation Loc, ExprResult Body) {
   auto Func = CurrentContextAsFunction();
   auto Type = Func->getType();
-  Body = TypecheckAssignment(Type, Body, Loc,
-                             getTokenRange(Func->getLocation()),
-                             Body.get()->getSourceRange());
-  CurrentContextAsFunction()->setBody(Body.get());
+  Body = CheckAndApplyAssignmentConstraints(Loc,
+                                            Type, Body.get(),
+                                            Sema::AssignmentAction::Returning);
+  if(Body.isUsable())
+    CurrentContextAsFunction()->setBody(Body.get());
 }
 
 void Sema::ActOnEndStatementFunction(ASTContext &C) {
diff --git a/lib/Sema/SemaDecl.cpp b/lib/Sema/SemaDecl.cpp
index 3cbb5c9fff..444bcb5034 100644
--- a/lib/Sema/SemaDecl.cpp
+++ b/lib/Sema/SemaDecl.cpp
@@ -261,7 +261,6 @@ Decl *Sema::ActOnParameterEntityDecl(ASTContext &C, QualType T,
     Value = CheckAndApplyAssignmentConstraints(EqualLoc,
                                                VD->getType(), Value.get(),
                                                Sema::AssignmentAction::Initializing);
-
     if(Value.isInvalid()) return nullptr;
     // FIXME: if value is invalid, mutate into parameter givin a zero value
     VD->MutateIntoParameter(Value.get());
diff --git a/lib/Sema/SemaExpr.cpp b/lib/Sema/SemaExpr.cpp
index 77cb9bf8cb..32ef2d1442 100644
--- a/lib/Sema/SemaExpr.cpp
+++ b/lib/Sema/SemaExpr.cpp
@@ -272,6 +272,9 @@ bool Sema::DiagnoseAssignmentResult(AssignConvertType ConvTy,
     case AssignmentAction::Passing:
       Diag = diag::err_typecheck_passing_incompatible;
       break;
+    case AssignmentAction::Returning:
+      Diag = diag::err_typecheck_return_incompatible;
+      break;
     default:
       llvm_unreachable("invalid assignment action");
       return true;
@@ -282,9 +285,10 @@ bool Sema::DiagnoseAssignmentResult(AssignConvertType ConvTy,
     if(DstExpr)
       Reporter << DstExpr->getSourceRange();
     if(Action.getType() == AssignmentAction::Assigning ||
-       Action.getType() == AssignmentAction::Initializing) {
+       Action.getType() == AssignmentAction::Initializing)
       Reporter << DstType << SrcType;
-    }
+    else if(Action.getType() == AssignmentAction::Returning)
+      Reporter << SrcType << DstType;
     break;
   }
   case IncompatibleDimensions:
diff --git a/test/Sema/statementFunctions.f95 b/test/Sema/statementFunctions.f95
index 41692ca606..f18ce023a8 100644
--- a/test/Sema/statementFunctions.f95
+++ b/test/Sema/statementFunctions.f95
@@ -16,7 +16,7 @@ PROGRAM test
   INTEGER FOO ! expected-note@+1 {{previous definition is here}}
   FOO(I,I) = 1 ! expected-error {{redefinition of 'i'}}
 
-  Z(A) = 'Hello' ! expected-error {{assigning to 'real' from incompatible type 'character'}}
+  Z(A) = 'Hello' ! expected-error {{returning 'character' from a function with incompatible result type 'real'}}
 
   I = X(2) ! CHECK: i = int(x(2))
   I = X() ! expected-error {{too few arguments to function call, expected 1, have 0}}