Added return type check for non-local return (despite it is prohibited yet)

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiUtil.java

@@ -1,5 +1,6 @@
 package org.jetbrains.jet.lang.psi;
 
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lexer.JetTokens;
@@ -62,4 +63,14 @@ public class JetPsiUtil {
         }
         return rootElements;
     }
+
+    @Nullable
+    public static JetNamedFunction getSurroundingFunction(@Nullable PsiElement element) {
+        while (element != null) {
+            if (element instanceof JetNamedFunction) return (JetNamedFunction) element;
+            if (element instanceof JetClassOrObject || element instanceof JetNamespace) return null;
+            element = element.getParent();
+        }
+        return null;
+    }
 }

compiler/frontend/src/org/jetbrains/jet/lang/resolve/BindingContextUtils.java

@@ -4,8 +4,12 @@ import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
+import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
 import org.jetbrains.jet.lang.descriptors.VariableDescriptor;
 import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.types.JetType;
+
+import static org.jetbrains.jet.lang.resolve.BindingContext.DECLARATION_TO_DESCRIPTOR;
 
 /**
  * @author abreslav
@@ -40,4 +44,12 @@ public class BindingContextUtils {
         }
         return null;
     }
+
+    @Nullable
+    public static JetType getFunctionReturnType(@NotNull BindingContext bindingContext, @NotNull JetFunction function) {
+        DeclarationDescriptor descriptor = bindingContext.get(DECLARATION_TO_DESCRIPTOR, function);
+        assert descriptor instanceof FunctionDescriptor;
+        FunctionDescriptor functionDescriptor = (FunctionDescriptor) descriptor;
+        return functionDescriptor.getReturnType();
+    }
 }

compiler/frontend/src/org/jetbrains/jet/lang/types/expressions/ControlStructureTypingVisitor.java

@@ -1,12 +1,14 @@
 package org.jetbrains.jet.lang.types.expressions;
 
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
 import org.jetbrains.jet.lang.descriptors.VariableDescriptor;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.resolve.BindingContextUtils;
 import org.jetbrains.jet.lang.resolve.calls.autocasts.DataFlowInfo;
 import org.jetbrains.jet.lang.resolve.calls.OverloadResolutionResults;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
@@ -46,7 +48,6 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
             JetType conditionType = facade.getType(condition, context.replaceScope(scope));
 
             if (conditionType != null && !isBoolean(context.semanticServices, conditionType)) {
-//                    context.trace.getErrorHandler().genericError(condition.getNode(), "Condition must be of type Boolean, but was of type " + conditionType);
                 context.trace.report(TYPE_MISMATCH_IN_CONDITION.on(condition, conditionType));
             }
         }
@@ -207,7 +208,6 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
                 if (expectedParameterType != null &&
                         actualParameterType != null &&
                         !context.semanticServices.getTypeChecker().isSubtypeOf(expectedParameterType, actualParameterType)) {
-//                        context.trace.getErrorHandler().genericError(typeReference.getNode(), "The loop iterates over values of type " + expectedParameterType + " but the parameter is declared to be " + actualParameterType);
                     context.trace.report(TYPE_MISMATCH_IN_FOR_LOOP.on(typeReference, expectedParameterType, actualParameterType));
                 }
             }
@@ -244,11 +244,9 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
             boolean hasNextPropertySupported = hasNextProperty != null;
             if (hasNextFunctionSupported && hasNextPropertySupported && !ErrorUtils.isErrorType(iteratorType)) {
                 // TODO : overload resolution rules impose priorities here???
-//                    context.trace.getErrorHandler().genericError(reportErrorsOn, "An ambiguity between 'iterator().hasNext()' function and 'iterator().hasNext' property");
                 context.trace.report(HAS_NEXT_PROPERTY_AND_FUNCTION_AMBIGUITY.on(loopRangeExpression));
             }
             else if (!hasNextFunctionSupported && !hasNextPropertySupported) {
-//                    context.trace.getErrorHandler().genericError(reportErrorsOn, "Loop range must have an 'iterator().hasNext()' function or an 'iterator().hasNext' property");
                 context.trace.report(HAS_NEXT_MISSING.on(loopRangeExpression));
             }
             else {
@@ -257,10 +255,8 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
 
             OverloadResolutionResults<FunctionDescriptor> nextResolutionResults = context.resolveExactSignature(new TransientReceiver(iteratorType), "next", Collections.<JetType>emptyList());
             if (nextResolutionResults.isAmbiguity()) {
-//                    context.trace.getErrorHandler().genericError(reportErrorsOn, "Method 'iterator().next()' is ambiguous for this expression");
                 context.trace.report(NEXT_AMBIGUITY.on(loopRangeExpression));
             } else if (nextResolutionResults.isNothing()) {
-//                    context.trace.getErrorHandler().genericError(reportErrorsOn, "Loop range must have an 'iterator().next()' method");
                 context.trace.report(NEXT_MISSING.on(loopRangeExpression));
             } else {
                 FunctionDescriptor nextFunction = nextResolutionResults.getResult().getResultingDescriptor();
@@ -278,7 +274,6 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
                 context.trace.report(ITERATOR_AMBIGUITY.on(loopRangeExpression, iteratorResolutionResults.getResults()));
             }
             else {
-//                    context.trace.getErrorHandler().genericError(reportErrorsOn, errorMessage);
                 context.trace.report(ITERATOR_MISSING.on(loopRangeExpression));
             }
         }
@@ -289,7 +284,6 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
     private FunctionDescriptor checkHasNextFunctionSupport(@NotNull JetExpression loopRange, @NotNull JetType iteratorType, ExpressionTypingContext context) {
         OverloadResolutionResults<FunctionDescriptor> hasNextResolutionResults = context.resolveExactSignature(new TransientReceiver(iteratorType), "hasNext", Collections.<JetType>emptyList());
         if (hasNextResolutionResults.isAmbiguity()) {
-//                context.trace.getErrorHandler().genericError(loopRange.getNode(), "Method 'iterator().hasNext()' is ambiguous for this expression");
             context.trace.report(HAS_NEXT_FUNCTION_AMBIGUITY.on(loopRange));
         } else if (hasNextResolutionResults.isNothing()) {
             return null;
@@ -297,7 +291,6 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
             assert hasNextResolutionResults.isSuccess();
             JetType hasNextReturnType = hasNextResolutionResults.getResult().getResultingDescriptor().getReturnType();
             if (!isBoolean(context.semanticServices, hasNextReturnType)) {
-//                    context.trace.getErrorHandler().genericError(loopRange.getNode(), "The 'iterator().hasNext()' method of the loop range must return Boolean, but returns " + hasNextReturnType);
                 context.trace.report(HAS_NEXT_FUNCTION_TYPE_MISMATCH.on(loopRange, hasNextReturnType));
             }
         }
@@ -314,11 +307,9 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
             JetType hasNextReturnType = hasNextProperty.getOutType();
             if (hasNextReturnType == null) {
                 // TODO : accessibility
-//                    context.trace.getErrorHandler().genericError(loopRange.getNode(), "The 'iterator().hasNext' property of the loop range must be readable");
                 context.trace.report(HAS_NEXT_MUST_BE_READABLE.on(loopRange));
             }
             else if (!isBoolean(context.semanticServices, hasNextReturnType)) {
-//                    context.trace.getErrorHandler().genericError(loopRange.getNode(), "The 'iterator().hasNext' property of the loop range must return Boolean, but returns " + hasNextReturnType);
                 context.trace.report(HAS_NEXT_PROPERTY_TYPE_MISMATCH.on(loopRange, hasNextReturnType));
             }
         }
@@ -379,20 +370,33 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
     public JetType visitReturnExpression(JetReturnExpression expression, ExpressionTypingContext context) {
         context.labelResolver.recordLabel(expression, context);
         if (context.expectedReturnType == TypeUtils.FORBIDDEN) {
-//                context.trace.getErrorHandler().genericError(expression.getNode(), "'return' is not allowed here");
             context.trace.report(RETURN_NOT_ALLOWED.on(expression));
             return null;
         }
         JetExpression returnedExpression = expression.getReturnedExpression();
 
-        JetType returnedType = JetStandardClasses.getUnitType();
+        JetType expectedType = context.expectedReturnType;
-        if (returnedExpression != null) {
+        if (expression.getTargetLabel() == null) {
-            facade.getType(returnedExpression, context.replaceExpectedType(context.expectedReturnType).replaceScope(context.scope));
+            if (PsiTreeUtil.getParentOfType(expression, JetFunctionLiteral.class) ==
+                PsiTreeUtil.getParentOfType(expression, JetDeclaration.class)) {  // expression is located in function literal
+                JetNamedFunction function = JetPsiUtil.getSurroundingFunction(expression);
+                if (function != null && function.getReturnTypeRef() != null) {
+                    expectedType = BindingContextUtils.getFunctionReturnType(context.trace.getBindingContext(), function);
+                }
+            }
         }
         else {
-            if (context.expectedReturnType != TypeUtils.NO_EXPECTED_TYPE && !JetStandardClasses.isUnit(context.expectedReturnType)) {
+            PsiElement element = context.trace.get(LABEL_TARGET, expression.getTargetLabel());
-//                    context.trace.getErrorHandler().genericError(expression.getNode(), "This function must return a value of type " + context.expectedReturnType);
+            if (element instanceof JetFunction && ((JetFunction) element).getReturnTypeRef() != null) {
-                context.trace.report(RETURN_TYPE_MISMATCH.on(expression, context.expectedReturnType));
+                expectedType = BindingContextUtils.getFunctionReturnType(context.trace.getBindingContext(), (JetFunction) element);
+            }
+        }
+        if (returnedExpression != null) {
+            facade.getType(returnedExpression, context.replaceExpectedType(expectedType).replaceScope(context.scope));
+        }
+        else {
+            if (expectedType != TypeUtils.NO_EXPECTED_TYPE && expectedType != null && !JetStandardClasses.isUnit(expectedType)) {
+                context.trace.report(RETURN_TYPE_MISMATCH.on(expression, expectedType));
             }
         }
         return DataFlowUtils.checkType(JetStandardClasses.getNothingType(), expression, context);

compiler/testData/checkerWithErrorTypes/full/FunctionReturnTypes.jet

@@ -5,11 +5,12 @@ fun unitEmpty() : Unit {}
 fun unitEmptyReturn() : Unit {return}
 fun unitIntReturn() : Unit {return <!TYPE_MISMATCH!>1<!>}
 fun unitUnitReturn() : Unit {return ()}
-fun test1() : Any = {<!RETURN_NOT_ALLOWED!>return<!>}
+fun test1() : Any = {<!RETURN_TYPE_MISMATCH, RETURN_NOT_ALLOWED!>return<!>}
 fun test2() : Any = @a {return@a 1}
 fun test3() : Any { <!RETURN_TYPE_MISMATCH!>return<!> }
-fun test4(): fun(): Unit = { <!RETURN_NOT_ALLOWED!>return@test4<!> }
+fun test4(): fun(): Unit = { <!RETURN_TYPE_MISMATCH, RETURN_NOT_ALLOWED!>return@test4<!> }
 fun test5(): Any = @{ return@ }
+fun test5(): Any = {<!RETURN_NOT_ALLOWED!>return 1<!>}
 
 fun bbb() {
   return <!TYPE_MISMATCH!>1<!>

compiler/testData/checkerWithErrorTypes/quick/regressions/kt411.jet

@@ -23,7 +23,7 @@ fun t2() : String {
         if (true) {
             return@ 1
         }
-        <!RETURN_NOT_ALLOWED!>return <!TYPE_MISMATCH!>"s"<!><!>
+        <!RETURN_NOT_ALLOWED!>return "s"<!>
     }
     return "s"
 }
@@ -32,10 +32,10 @@ fun t3() : String {
     invoker(
     @{
         if (true) {
-            <!RETURN_NOT_ALLOWED!>return@t3 <!TYPE_MISMATCH!>"1"<!><!>
+            <!RETURN_NOT_ALLOWED!>return@t3 "1"<!>
         }
         else {
-            <!RETURN_NOT_ALLOWED!>return <!TYPE_MISMATCH!>"2"<!><!>
+            <!RETURN_NOT_ALLOWED!>return <!ERROR_COMPILE_TIME_VALUE!>2<!><!>
         }
         return@ 0
     }
@@ -47,7 +47,7 @@ fun t3() : String {
     )
     invoker(
     {
-        <!RETURN_NOT_ALLOWED!>return <!TYPE_MISMATCH!>"0"<!><!>
+        <!RETURN_NOT_ALLOWED!>return "0"<!>
     }
     )
     return "2"