KT-411 Wrong type expected when returning from a function literal

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowBuilder.java

@@ -49,8 +49,10 @@ public interface JetControlFlowBuilder {
 
     void exitSubroutine(@NotNull JetDeclaration subroutine);
 
-    @Nullable
+    @NotNull
     JetElement getCurrentSubroutine();
+    @Nullable
+    JetElement getReturnSubroutine();
     void returnValue(@NotNull JetExpression returnExpression, @NotNull JetElement subroutine);
 
     void returnNoValue(@NotNull JetElement returnExpression, @NotNull JetElement subroutine);

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowBuilderAdapter.java

@@ -140,13 +140,20 @@ public class JetControlFlowBuilderAdapter implements JetControlFlowBuilder {
         builder.exitSubroutine(subroutine);
     }
 
+    @NotNull
     @Override
-    @Nullable
     public JetElement getCurrentSubroutine() {
         assert builder != null;
         return builder.getCurrentSubroutine();
     }
 
+    @Override
+    @Nullable
+    public JetElement getReturnSubroutine() {
+        assert builder != null;
+        return builder.getReturnSubroutine();
+    }
+
     @Override
     public void returnValue(@NotNull JetExpression returnExpression, @NotNull JetElement subroutine) {
         assert builder != null;

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetControlFlowProcessor.java

@@ -528,8 +528,8 @@ public class JetControlFlowProcessor {
             }
             JetSimpleNameExpression labelElement = expression.getTargetLabel();
             JetElement subroutine;
+            String labelName = expression.getLabelName();
             if (labelElement != null) {
-                String labelName = expression.getLabelName();
                 assert labelName != null;
                 PsiElement labeledElement = BindingContextUtils.resolveToDeclarationPsiElement(trace.getBindingContext(), labelElement);
                 if (labeledElement != null) {
@@ -541,13 +541,18 @@ public class JetControlFlowProcessor {
                 }
             }
             else {
-                subroutine = builder.getCurrentSubroutine();
+                subroutine = builder.getReturnSubroutine();
                 // TODO : a context check
             }
             //todo cache JetFunctionLiteral instead
             if (subroutine instanceof JetFunctionLiteralExpression) {
                 subroutine = ((JetFunctionLiteralExpression) subroutine).getFunctionLiteral();
             }
+            boolean error = false;
+            if (builder.getCurrentSubroutine() != subroutine) {
+                trace.report(RETURN_NOT_ALLOWED.on(expression));
+                error = true;
+            }
             if (subroutine instanceof JetFunction || subroutine instanceof JetPropertyAccessor || subroutine instanceof JetConstructor) {
                 if (returnedExpression == null) {
                     builder.returnNoValue(expression, subroutine);
@@ -556,9 +561,9 @@ public class JetControlFlowProcessor {
                     builder.returnValue(expression, subroutine);
                 }
             }
-            else {
+            else if (!error) {
                 if (labelElement != null) {
-                    trace.report(NOT_A_RETURN_LABEL.on(expression, expression.getLabelName()));
+                    trace.report(NOT_A_RETURN_LABEL.on(expression, labelName));
                 }
                 else {
                     trace.report(RETURN_NOT_ALLOWED.on(expression));

compiler/frontend/src/org/jetbrains/jet/lang/cfg/JetFlowInformationProvider.java

@@ -75,17 +75,22 @@ public class JetFlowInformationProvider {
         Pseudocode pseudocode = pseudocodeMap.get(subroutine);
         assert pseudocode != null;
 
+        final Set<Instruction> instructions = Sets.newHashSet(pseudocode.getInstructions());
         SubroutineExitInstruction exitInstruction = pseudocode.getExitInstruction();
         for (Instruction previousInstruction : exitInstruction.getPreviousInstructions()) {
             previousInstruction.accept(new InstructionVisitor() {
                 @Override
                 public void visitReturnValue(ReturnValueInstruction instruction) {
-                    returnedExpressions.add((JetExpression) instruction.getElement());
+                    if (instructions.contains(instruction)) { //exclude non-local return expressions
+                        returnedExpressions.add((JetExpression) instruction.getElement());
+                    }
                 }
 
                 @Override
                 public void visitReturnNoValue(ReturnNoValueInstruction instruction) {
-                    returnedExpressions.add((JetExpression) instruction.getElement());
+                    if (instructions.contains(instruction)) {
+                        returnedExpressions.add((JetExpression) instruction.getElement());
+                    }
                 }
 
 

compiler/frontend/src/org/jetbrains/jet/lang/cfg/pseudocode/JetControlFlowInstructionsGenerator.java

@@ -45,7 +45,12 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
 
     @Override
     public void enterSubroutine(@NotNull JetDeclaration subroutine) {
-        pushBuilder(subroutine, subroutine);
+        if (builder != null && subroutine instanceof JetFunctionLiteral) {
+            pushBuilder(subroutine, builder.getReturnSubroutine());
+        }
+        else {
+            pushBuilder(subroutine, subroutine);
+        }
         assert builder != null;
         builder.enterSubroutine(subroutine);
     }
@@ -66,13 +71,13 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
         private final Pseudocode pseudocode;
         private final Label error;
         private final Label sink;
-        private final JetElement currentSubroutine;
+        private final JetElement returnSubroutine;
 
-        private JetControlFlowInstructionsGeneratorWorker(@NotNull JetElement scopingElement, @NotNull JetElement currentSubroutine) {
+        private JetControlFlowInstructionsGeneratorWorker(@NotNull JetElement scopingElement, @NotNull JetElement returnSubroutine) {
             this.pseudocode = new Pseudocode(scopingElement);
             this.error = pseudocode.createLabel("error");
             this.sink = pseudocode.createLabel("sink");
-            this.currentSubroutine = currentSubroutine;
+            this.returnSubroutine = returnSubroutine;
         }
 
         public Pseudocode getPseudocode() {
@@ -134,9 +139,15 @@ public class JetControlFlowInstructionsGenerator extends JetControlFlowBuilderAd
             add(new SubroutineEnterInstruction(subroutine));
         }
 
+        @NotNull
         @Override
         public JetElement getCurrentSubroutine() {
-            return currentSubroutine;// subroutineInfo.empty() ? null : subroutineInfo.peek().getElement();
+            return pseudocode.getCorrespondingElement();
+        }
+
+        @Override
+        public JetElement getReturnSubroutine() {
+            return returnSubroutine;// subroutineInfo.empty() ? null : subroutineInfo.peek().getElement();
         }
 
         @Override

compiler/frontend/src/org/jetbrains/jet/lang/descriptors/FunctionDescriptorUtil.java

@@ -89,7 +89,7 @@ public class FunctionDescriptorUtil {
                                       ReceiverDescriptor.NO_RECEIVER,
                                       Collections.<TypeParameterDescriptor>emptyList(),
                                       JetStandardClasses.getValueParameters(functionDescriptor, functionType),
-                                      JetStandardClasses.getReturnType(functionType),
+                                      JetStandardClasses.obtainReturnTypeFromFunctionType(functionType),
                                       Modality.FINAL, Visibility.LOCAL);
     }
 }

compiler/frontend/src/org/jetbrains/jet/lang/types/JetStandardClasses.java

@@ -385,7 +385,7 @@ public class JetStandardClasses {
     }
 
     @NotNull
-    public static JetType getReturnType(@NotNull JetType type) {
+    public static JetType obtainReturnTypeFromFunctionType(@NotNull JetType type) {
         assert isFunctionType(type);
         List<TypeProjection> arguments = type.getArguments();
         return arguments.get(arguments.size() - 1).getType();

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

@@ -61,6 +61,7 @@ public class ClosureExpressionsTypingVisitor extends ExpressionTypingVisitor {
     @Override
     public JetType visitFunctionLiteralExpression(JetFunctionLiteralExpression expression, ExpressionTypingContext context) {
         JetFunctionLiteral functionLiteral = expression.getFunctionLiteral();
+        if (functionLiteral.getBodyExpression() == null) return null;
 
         JetTypeReference receiverTypeRef = functionLiteral.getReceiverTypeRef();
         final JetType receiverType;
@@ -137,19 +138,22 @@ public class ClosureExpressionsTypingVisitor extends ExpressionTypingVisitor {
         JetTypeReference returnTypeRef = functionLiteral.getReturnTypeRef();
         if (returnTypeRef != null) {
             returnType = context.getTypeResolver().resolveType(context.scope, returnTypeRef);
-            context.getServices().checkFunctionReturnType(expression, context.replaceScope(functionInnerScope).replaceExpectedReturnType(returnType).replaceDataFlowInfo(context.dataFlowInfo));
+            context.getServices().checkFunctionReturnType(expression, context.replaceScope(functionInnerScope).
+                    replaceExpectedType(returnType).replaceExpectedReturnType(returnType).replaceDataFlowInfo(context.dataFlowInfo));
         }
         else {
             if (functionTypeExpected) {
-                returnType = JetStandardClasses.getReturnType(expectedType);
+                returnType = JetStandardClasses.obtainReturnTypeFromFunctionType(expectedType);
             }
-            returnType = context.getServices().getBlockReturnedType(functionInnerScope, functionLiteral.getBodyExpression(), CoercionStrategy.COERCION_TO_UNIT, context.replaceExpectedType(returnType).replaceExpectedReturnType(returnType));
+//        }
+            returnType = context.getServices().getBlockReturnedType(functionInnerScope, functionLiteral.getBodyExpression(), CoercionStrategy.COERCION_TO_UNIT,
+                                                                    context.replaceExpectedType(returnType).replaceExpectedReturnType(returnType));
         }
         JetType safeReturnType = returnType == null ? ErrorUtils.createErrorType("<return type>") : returnType;
         functionDescriptor.setReturnType(safeReturnType);
 
         if (functionTypeExpected) {
-            JetType expectedReturnType = JetStandardClasses.getReturnType(expectedType);
+            JetType expectedReturnType = JetStandardClasses.obtainReturnTypeFromFunctionType(expectedType);
             if (JetStandardClasses.isUnit(expectedReturnType)) {
                 functionDescriptor.setReturnType(expectedReturnType);
                 return DataFlowUtils.checkType(JetStandardClasses.getFunctionType(Collections.<AnnotationDescriptor>emptyList(), effectiveReceiverType, parameterTypes, expectedReturnType), expression, context);

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

@@ -122,7 +122,9 @@ public class ExpressionTypingServices {
 
         if (function instanceof JetFunctionLiteralExpression) {
             JetFunctionLiteralExpression functionLiteralExpression = (JetFunctionLiteralExpression) function;
-            getBlockReturnedType(newContext.scope, functionLiteralExpression.getBodyExpression(), CoercionStrategy.COERCION_TO_UNIT, newContext);
+            JetBlockExpression blockExpression = functionLiteralExpression.getBodyExpression();
+            assert blockExpression != null;
+            getBlockReturnedType(newContext.scope, blockExpression, CoercionStrategy.COERCION_TO_UNIT, context);
         }
         else {
             expressionTypingFacade.getType(bodyExpression, newContext);

compiler/testData/checkerWithErrorTypes/full/FunctionReturnTypes.jet

@@ -5,9 +5,11 @@ fun unitEmpty() : Unit {}
 fun unitEmptyReturn() : Unit {return}
 fun unitIntReturn() : Unit {return <!TYPE_MISMATCH!>1<!>}
 fun unitUnitReturn() : Unit {return ()}
-fun test1() : Any = {return}
+fun test1() : Any = {<!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 test5(): Any = @{ return@ }
 
 fun bbb() {
   return <!TYPE_MISMATCH!>1<!>

compiler/testData/checkerWithErrorTypes/full/Return.jet

@@ -6,12 +6,12 @@ class A {
             if (1 < 2)
                 return@inner
             else
-                return@outer
+                <!RETURN_NOT_ALLOWED!>return@outer<!>
         }
         if (1 < 2)
-            <!NOT_A_RETURN_LABEL!>return@A<!>
+            <!RETURN_NOT_ALLOWED!>return@A<!>
         else if (2 < 3)
-            <!NOT_A_RETURN_LABEL!>return<!UNRESOLVED_REFERENCE!>@inner<!><!>
+            <!RETURN_NOT_ALLOWED!>return<!UNRESOLVED_REFERENCE!>@inner<!><!>
         return@outer
     }
 }

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

@@ -0,0 +1,69 @@
+//kt-411 Wrong type expected when returning from a function literal
+
+namespace kt411
+
+fun f() {
+    invoker(
+    @{
+        return@ 11 // expects Function, but should expect Int
+    }
+    )
+}
+fun invoker(gen : fun() : Int) : Int  = 0
+
+//more tests
+fun t1() {
+    val v = @{ () : Int =>
+        return@ 111
+    }
+}
+
+fun t2() : String {
+    val g : fun(): Int = @{
+        if (true) {
+            return@ 1
+        }
+        <!RETURN_NOT_ALLOWED!>return <!TYPE_MISMATCH!>"s"<!><!>
+    }
+    return "s"
+}
+
+fun t3() : String {
+    invoker(
+    @{
+        if (true) {
+            <!RETURN_NOT_ALLOWED!>return@t3 <!TYPE_MISMATCH!>"1"<!><!>
+        }
+        else {
+            <!RETURN_NOT_ALLOWED!>return <!TYPE_MISMATCH!>"2"<!><!>
+        }
+        return@ 0
+    }
+    )
+    invoker(
+    @{ (): Int =>
+        return@ 1
+    }
+    )
+    invoker(
+    {
+        <!RETURN_NOT_ALLOWED!>return <!TYPE_MISMATCH!>"0"<!><!>
+    }
+    )
+    return "2"
+}
+
+fun t4() : Int {
+    val h : fun (): String = @l{
+        return@l "a"
+    }
+    val g : fun (): String = @{ () : String =>
+        return@ "a"
+    }
+
+    fun inner(): String {
+        return "2"
+    }
+
+    return 12
+}

idea/testData/checker/FunctionReturnTypes.jet

@@ -5,7 +5,7 @@ fun unitEmpty() : Unit {}
 fun unitEmptyReturn() : Unit {return}
 fun unitIntReturn() : Unit {return <error>1</error>}
 fun unitUnitReturn() : Unit {return ()}
-fun test1() : Any = {return}
+fun test1() : Any = { <error>return</error> }
 fun test2() : Any = @a {return@a 1}
 fun test3() : Any { <error>return</error> }