Support local returns in lambdas

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

@@ -18,6 +18,8 @@ package org.jetbrains.jet.lang.types.expressions;
 
 import com.google.common.collect.Lists;
 import com.intellij.psi.PsiElement;
+import com.intellij.util.Function;
+import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.*;
@@ -28,21 +30,20 @@ import org.jetbrains.jet.lang.resolve.*;
 import org.jetbrains.jet.lang.resolve.calls.CallResolverUtil;
 import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
-import org.jetbrains.jet.lang.types.DeferredType;
+import org.jetbrains.jet.lang.types.*;
-import org.jetbrains.jet.lang.types.ErrorUtils;
-import org.jetbrains.jet.lang.types.JetType;
-import org.jetbrains.jet.lang.types.JetTypeInfo;
 import org.jetbrains.jet.lang.types.checker.JetTypeChecker;
 import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
 import org.jetbrains.jet.util.lazy.RecursionIntolerantLazyValueWithDefault;
 import org.jetbrains.jet.util.slicedmap.WritableSlice;
 
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 
 import static org.jetbrains.jet.lang.diagnostics.Errors.*;
 import static org.jetbrains.jet.lang.resolve.BindingContext.*;
 import static org.jetbrains.jet.lang.types.TypeUtils.NO_EXPECTED_TYPE;
+import static org.jetbrains.jet.lang.types.expressions.CoercionStrategy.COERCION_TO_UNIT;
 import static org.jetbrains.jet.lang.types.expressions.ExpressionTypingUtils.CANT_INFER_LAMBDA_PARAM_TYPE;
 
 public class ClosureExpressionsTypingVisitor extends ExpressionTypingVisitor {
@@ -128,7 +129,8 @@ public class ClosureExpressionsTypingVisitor extends ExpressionTypingVisitor {
         AnonymousFunctionDescriptor functionDescriptor = new AnonymousFunctionDescriptor(
                 context.scope.getContainingDeclaration(), Collections.<AnnotationDescriptor>emptyList(), CallableMemberDescriptor.Kind.DECLARATION);
 
-        List<ValueParameterDescriptor> valueParameterDescriptors = createValueParameterDescriptors(context, functionLiteral, functionDescriptor, functionTypeExpected);
+        List<ValueParameterDescriptor> valueParameterDescriptors = createValueParameterDescriptors(context, functionLiteral,
+                                                                                                   functionDescriptor, functionTypeExpected);
 
         JetType effectiveReceiverType;
         if (receiverTypeRef == null) {
@@ -243,7 +245,8 @@ public class ClosureExpressionsTypingVisitor extends ExpressionTypingVisitor {
             @NotNull SimpleFunctionDescriptorImpl functionDescriptor,
             boolean functionTypeExpected
     ) {
-        TemporaryBindingTrace temporaryTrace = TemporaryBindingTrace.create(context.trace, "trace to resolve function literal expression", expression);
+        TemporaryBindingTrace temporaryTrace = TemporaryBindingTrace.create(context.trace, "trace to resolve function literal expression",
+                                                                            expression);
         JetType expectedReturnType = functionTypeExpected ? KotlinBuiltIns.getInstance().getReturnTypeFromFunctionType(context.expectedType) : null;
         JetType returnType = computeUnsafeReturnType(expression, context, functionDescriptor, temporaryTrace, expectedReturnType);
 
@@ -276,19 +279,79 @@ public class ClosureExpressionsTypingVisitor extends ExpressionTypingVisitor {
 
         JetScope functionInnerScope = FunctionDescriptorUtil.getFunctionInnerScope(context.scope, functionDescriptor, context.trace);
         JetTypeReference returnTypeRef = functionLiteral.getReturnTypeRef();
+        JetType declaredReturnType = null;
         if (returnTypeRef != null) {
-            JetType returnType = context.expressionTypingServices.getTypeResolver().resolveType(context.scope, returnTypeRef, context.trace, true);
+            declaredReturnType = context.expressionTypingServices.getTypeResolver().resolveType(context.scope, returnTypeRef, context.trace, true);
-            context.expressionTypingServices.checkFunctionReturnType(expression.getFunctionLiteral(), context.replaceScope(functionInnerScope).
-                    replaceExpectedType(returnType).replaceBindingTrace(temporaryTrace), temporaryTrace);
             if (expectedReturnType != null) {
-                if (!JetTypeChecker.INSTANCE.isSubtypeOf(returnType, expectedReturnType)) {
+                if (!JetTypeChecker.INSTANCE.isSubtypeOf(declaredReturnType, expectedReturnType)) {
                     temporaryTrace.report(EXPECTED_RETURN_TYPE_MISMATCH.on(returnTypeRef, expectedReturnType));
                 }
             }
-            return returnType;
         }
-        ExpressionTypingContext newContext = context.replaceExpectedType(expectedReturnType != null ? expectedReturnType : NO_EXPECTED_TYPE)
+
-                .replaceBindingTrace(temporaryTrace).replaceScope(functionInnerScope);
+        // Type-check the body
-        return context.expressionTypingServices.getBlockReturnedType(bodyExpression, CoercionStrategy.COERCION_TO_UNIT, newContext).getType();
+        ExpressionTypingContext newContext = context.replaceBindingTrace(temporaryTrace).replaceScope(functionInnerScope)
+                .replaceExpectedType(declaredReturnType != null
+                                     ? declaredReturnType
+                                     : (expectedReturnType != null ? expectedReturnType : NO_EXPECTED_TYPE));
+
+        JetType typeOfBodyExpression = context.expressionTypingServices.getBlockReturnedType(bodyExpression, COERCION_TO_UNIT, newContext).getType();
+
+        List<JetType> returnedExpressionTypes = Lists.newArrayList(getTypesOfLocallyReturnedExpressions(
+                functionLiteral, temporaryTrace, collectReturns(bodyExpression)));
+        ContainerUtil.addIfNotNull(returnedExpressionTypes, typeOfBodyExpression);
+
+        if (declaredReturnType != null) return declaredReturnType;
+        if (returnedExpressionTypes.isEmpty()) return null;
+        return CommonSupertypes.commonSupertype(returnedExpressionTypes);
+    }
+
+    private static List<JetType> getTypesOfLocallyReturnedExpressions(
+            final JetFunctionLiteral functionLiteral,
+            final TemporaryBindingTrace temporaryTrace,
+            Collection<JetReturnExpression> returnExpressions
+    ) {
+        return ContainerUtil.mapNotNull(returnExpressions, new Function<JetReturnExpression, JetType>() {
+            @Override
+            public JetType fun(JetReturnExpression returnExpression) {
+                JetSimpleNameExpression label = returnExpression.getTargetLabel();
+                if (label == null) {
+                    // No label => non-local return
+                    return null;
+                }
+
+                PsiElement labelTarget = temporaryTrace.get(BindingContext.LABEL_TARGET, label);
+                if (labelTarget != functionLiteral) {
+                    // Either a local return of inner lambda/function or a non-local return
+                    return null;
+                }
+
+                JetExpression returnedExpression = returnExpression.getReturnedExpression();
+                if (returnedExpression == null) {
+                    return KotlinBuiltIns.getInstance().getUnitType();
+                }
+                JetType returnedType = temporaryTrace.get(EXPRESSION_TYPE, returnedExpression);
+                assert returnedType != null : "No type for returned expression: " + returnedExpression + ",\n" +
+                                              "the type should have been computed by getBlockReturnedType() above";
+                return returnedType;
+            }
+        });
+    }
+
+    public static Collection<JetReturnExpression> collectReturns(@NotNull JetExpression expression) {
+        Collection<JetReturnExpression> result = Lists.newArrayList();
+        expression.accept(
+                new JetTreeVisitor<Collection<JetReturnExpression>>() {
+                    @Override
+                    public Void visitReturnExpression(
+                            JetReturnExpression expression, Collection<JetReturnExpression> data
+                    ) {
+                        data.add(expression);
+                        return null;
+                    }
+                },
+                result
+        );
+        return result;
     }
 }

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

@@ -154,10 +154,10 @@ public class ExpressionTypingServices {
         }
         checkFunctionReturnType(function, ExpressionTypingContext.newContext(
                 this, trace, functionInnerScope, dataFlowInfo, expectedReturnType != null ? expectedReturnType : NO_EXPECTED_TYPE, ExpressionPosition.FREE
-        ), trace);
+        ));
     }
 
-    /*package*/ void checkFunctionReturnType(JetDeclarationWithBody function, ExpressionTypingContext context, BindingTrace trace) {
+    /*package*/ void checkFunctionReturnType(JetDeclarationWithBody function, ExpressionTypingContext context) {
         JetExpression bodyExpression = function.getBodyExpression();
         if (bodyExpression == null) return;
 
@@ -167,15 +167,7 @@ public class ExpressionTypingServices {
                 ? context.replaceExpectedType(NO_EXPECTED_TYPE)
                 : context;
 
-        if (function instanceof JetFunctionLiteral) {
+        expressionTypingFacade.getTypeInfo(bodyExpression, newContext, !blockBody);
-            JetFunctionLiteral functionLiteral = (JetFunctionLiteral) function;
-            JetBlockExpression blockExpression = functionLiteral.getBodyExpression();
-            assert blockExpression != null;
-            getBlockReturnedType(blockExpression, CoercionStrategy.COERCION_TO_UNIT, context);
-        }
-        else {
-            expressionTypingFacade.getTypeInfo(bodyExpression, newContext, !blockBody);
-        }
     }
 
     @NotNull

compiler/testData/diagnostics/tests/functionLiterals/return/LocalAndNonLocalReturnInLambda.kt

@@ -0,0 +1,8 @@
+fun test2(a: Int) {
+    (run @f{
+      if (a > 0) <!RETURN_NOT_ALLOWED!>return<!>
+      return@f 1
+    }): Int
+}
+
+fun run<T>(f: () -> T): T { return f() }

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnExplicitLabelNoParens.kt

@@ -1,5 +1,5 @@
 fun test2() {
-    run @f{return@f 1}
+    (run @f{return@f 1}): Int
 }
 
 fun run<T>(f: () -> T): T { return f() }

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnExplicitLabelParens.kt

@@ -1,10 +1,10 @@
 fun test() {
-    run(@f{return@f 1})
+    run(@f{return@f 1}): Int
 }
 
 
 fun test1() {
-    run(@{return@ 1})
+    run(@{return@ 1}): Int
 }
 
 fun run<T>(f: () -> T): T { return f() }

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnInNestedFunction.kt

@@ -0,0 +1,11 @@
+fun test() {
+    (run @f{
+      fun local(a: Int): String {
+        if (a > 0) return "2"
+        return@local "3"
+      }
+      return@f 1
+    }): Int
+}
+
+fun run<T>(f: () -> T): T { return f() }

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnInNestedLambda.kt

@@ -0,0 +1,10 @@
+fun test() {
+    (run @f{
+      run @ff {
+        return@ff "2"
+      }
+      return@f 1
+    }): Int
+}
+
+fun run<T>(f: () -> T): T { return f() }

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnNoCoercionToUnit.kt

@@ -0,0 +1,8 @@
+fun test(a: Int) {
+    run @f{
+      if (a > 0) return@f
+      else return@f 1
+    }
+}
+
+fun run<T>(f: () -> T): T { return f() }

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnUnit.kt

@@ -0,0 +1,8 @@
+fun test(a: Int) {
+    (run @f{
+      if (a > 0) return@f
+      else return@f Unit.VALUE
+    }): Unit
+}
+
+fun run<T>(f: () -> T): T { return f() }

compiler/testData/diagnostics/tests/functionLiterals/return/MixedReturnsFromLambda.kt

@@ -0,0 +1,13 @@
+trait A
+trait B: A
+trait C: A
+
+
+fun test(a: C, b: B) {
+    (run @f{
+      if (a != b) return@f a
+      b
+    }): A
+}
+
+fun run<T>(f: () -> T): T { return f() }

compiler/tests/org/jetbrains/jet/checkers/JetDiagnosticsTestGenerated.java

@@ -2437,6 +2437,11 @@ public class JetDiagnosticsTestGenerated extends AbstractDiagnosticsTestWithEage
                     doTest("compiler/testData/diagnostics/tests/functionLiterals/return/ForbiddenNonLocalReturnNoType.kt");
                 }
                 
+                @TestMetadata("LocalAndNonLocalReturnInLambda.kt")
+                public void testLocalAndNonLocalReturnInLambda() throws Exception {
+                    doTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalAndNonLocalReturnInLambda.kt");
+                }
+                
                 @TestMetadata("LocalReturnExplicitLabelNoParens.kt")
                 public void testLocalReturnExplicitLabelNoParens() throws Exception {
                     doTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnExplicitLabelNoParens.kt");
@@ -2447,6 +2452,31 @@ public class JetDiagnosticsTestGenerated extends AbstractDiagnosticsTestWithEage
                     doTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnExplicitLabelParens.kt");
                 }
                 
+                @TestMetadata("LocalReturnInNestedFunction.kt")
+                public void testLocalReturnInNestedFunction() throws Exception {
+                    doTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnInNestedFunction.kt");
+                }
+                
+                @TestMetadata("LocalReturnInNestedLambda.kt")
+                public void testLocalReturnInNestedLambda() throws Exception {
+                    doTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnInNestedLambda.kt");
+                }
+                
+                @TestMetadata("LocalReturnNoCoercionToUnit.kt")
+                public void testLocalReturnNoCoercionToUnit() throws Exception {
+                    doTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnNoCoercionToUnit.kt");
+                }
+                
+                @TestMetadata("LocalReturnUnit.kt")
+                public void testLocalReturnUnit() throws Exception {
+                    doTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnUnit.kt");
+                }
+                
+                @TestMetadata("MixedReturnsFromLambda.kt")
+                public void testMixedReturnsFromLambda() throws Exception {
+                    doTest("compiler/testData/diagnostics/tests/functionLiterals/return/MixedReturnsFromLambda.kt");
+                }
+                
             }
             
             public static Test innerSuite() {