Added checking for ONLY_LOCAL_RETURN in actual parameter of inline function

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

@@ -23,7 +23,11 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.bindingContextUtil.BindingContextUtilPackage;
+import org.jetbrains.jet.lang.resolve.calls.model.ArgumentMapping;
+import org.jetbrains.jet.lang.resolve.calls.model.ArgumentMatch;
 import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCall;
+import org.jetbrains.jet.lang.resolve.calls.util.CallMaker;
+import org.jetbrains.jet.lang.types.lang.InlineUtil;
 
 public class InlineDescriptorUtils {
 
@@ -43,16 +47,23 @@ public class InlineDescriptorUtils {
         while (containingFunction instanceof JetFunctionLiteral && fromFunction != containingFunctionDescriptor) {
             //JetFunctionLiteralExpression
             containingFunction = containingFunction.getParent();
-            boolean isInlinedLambda = false;
+            boolean allowsNonLocalReturns = false;
             JetExpression call = JetPsiUtil.getParentCallIfPresent((JetFunctionLiteralExpression) containingFunction);
             if (call != null) {
                 ResolvedCall<?> resolvedCall = BindingContextUtilPackage.getResolvedCall(call, bindingContext);
                 CallableDescriptor resultingDescriptor = resolvedCall == null ? null : resolvedCall.getResultingDescriptor();
-                if (resultingDescriptor instanceof SimpleFunctionDescriptor) {
+                if (resultingDescriptor instanceof SimpleFunctionDescriptor &&
-                    isInlinedLambda = ((SimpleFunctionDescriptor) resultingDescriptor).getInlineStrategy().isInline();
+                    ((SimpleFunctionDescriptor) resultingDescriptor).getInlineStrategy().isInline()) {
+                    ValueArgument argument = getContainingArgument(containingFunction, call);
+                    if (argument != null) {
+                        ArgumentMapping mapping = resolvedCall.getArgumentMapping(argument);
+                        if (mapping instanceof ArgumentMatch) {
+                            allowsNonLocalReturns = allowsNonLocalReturns(((ArgumentMatch) mapping).getValueParameter());
+                        }
+                    }
                 }
             }
-            if (!isInlinedLambda) {
+            if (!allowsNonLocalReturns) {
                 return false;
             }
 
@@ -66,6 +77,22 @@ public class InlineDescriptorUtils {
         return fromFunction == containingFunctionDescriptor;
     }
 
+    @Nullable
+    public static ValueArgument getContainingArgument(PsiElement expression, @NotNull JetExpression stopAtCall) {
+        if (expression instanceof JetValueArgument) {
+            return (JetValueArgument) expression;
+        }
+
+        while (expression != null && expression.getParent() != stopAtCall) {
+            PsiElement parent = expression.getParent();
+            if (parent instanceof JetValueArgument) {
+                return (JetValueArgument) parent;
+            }
+            expression = parent;
+        }
+        return expression != null ? CallMaker.makeValueArgument((JetExpression) expression) : null;
+    }
+
     @Nullable
     public static DeclarationDescriptor getContainingClassOrFunctionDescriptor(@NotNull DeclarationDescriptor descriptor, boolean strict) {
         DeclarationDescriptor currentDescriptor = strict ? descriptor.getContainingDeclaration() : descriptor;
@@ -78,4 +105,14 @@ public class InlineDescriptorUtils {
 
         return null;
     }
+
+    public static boolean allowsNonLocalReturns(@NotNull CallableDescriptor lambdaDescriptor) {
+        if (lambdaDescriptor instanceof ValueParameterDescriptor) {
+            if (InlineUtil.hasOnlyLocalReturn((ValueParameterDescriptor) lambdaDescriptor)) {
+                //annotated
+                return false;
+            }
+        }
+        return true;
+    }
 }

compiler/frontend/src/org/jetbrains/jet/lang/resolve/calls/InlineCallResolverExtension.java

@@ -24,7 +24,6 @@ import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.diagnostics.Errors;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.DescriptorUtils;
-import org.jetbrains.jet.lang.resolve.InlineDescriptorUtils;
 import org.jetbrains.jet.lang.resolve.bindingContextUtil.BindingContextUtilPackage;
 import org.jetbrains.jet.lang.resolve.calls.context.BasicCallResolutionContext;
 import org.jetbrains.jet.lang.resolve.calls.model.DefaultValueArgument;
@@ -44,6 +43,9 @@ import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 
+import static org.jetbrains.jet.lang.resolve.InlineDescriptorUtils.allowsNonLocalReturns;
+import static org.jetbrains.jet.lang.resolve.InlineDescriptorUtils.checkNonLocalReturnUsage;
+
 public class InlineCallResolverExtension implements CallResolverExtension {
 
     private final SimpleFunctionDescriptor descriptor;
@@ -263,7 +265,7 @@ public class InlineCallResolverExtension implements CallResolverExtension {
     ) {
         if (!allowsNonLocalReturns(inlinableParameterDescriptor)) return;
 
-        if (!InlineDescriptorUtils.checkNonLocalReturnUsage(descriptor, parameterUsage, context.trace)) {
+        if (!checkNonLocalReturnUsage(descriptor, parameterUsage, context.trace)) {
             context.trace.report(Errors.NON_LOCAL_RETURN_NOT_ALLOWED.on(parameterUsage, parameterUsage, inlinableParameterDescriptor, descriptor));
         }
     }
@@ -276,13 +278,4 @@ public class InlineCallResolverExtension implements CallResolverExtension {
         return expression;
     }
 
-    private static boolean allowsNonLocalReturns(CallableDescriptor lambdaDescriptor) {
-        if (lambdaDescriptor instanceof ValueParameterDescriptor) {
-            if (InlineUtil.hasOnlyLocalReturn((ValueParameterDescriptor) lambdaDescriptor)) {
-                //annotated
-                return false;
-            }
-        }
-        return true;
-    }
 }

compiler/testData/codegen/boxInline/lambdaTransformation/sameCaptured.1.kt

@@ -1,4 +1,5 @@
 import test.*
+import kotlin.InlineOption.*
 
 fun testSameCaptured() : String {
     var result = 0;
@@ -6,7 +7,7 @@ fun testSameCaptured() : String {
     return if (result == 12) "OK" else "fail ${result}"
 }
 
-inline fun testSameCaptured(lambdaWithResultCaptured: () -> Unit) : String {
+inline fun testSameCaptured(inlineOptions(ONLY_LOCAL_RETURN) lambdaWithResultCaptured: () -> Unit) : String {
     var result = 1;
     result = doWork({result+=11; lambdaWithResultCaptured(); result})
     return if (result == 12) "OK" else "fail ${result}"

compiler/testData/diagnostics/tests/inline/nonLocalReturns/lambdaWithGlobalReturnsInsideOnlyLocalOne.kt

@@ -0,0 +1,10 @@
+import kotlin.InlineOption.ONLY_LOCAL_RETURN
+
+inline fun testSameCaptured(lambdaWithResultCaptured: () -> Unit) : String {
+    doWork({<!NON_LOCAL_RETURN_NOT_ALLOWED!>lambdaWithResultCaptured<!>()})
+    return "OK"
+}
+
+inline fun <R> doWork(inlineOptions(ONLY_LOCAL_RETURN) job: ()-> R) : R {
+    return job()
+}

compiler/testData/diagnostics/tests/inline/nonLocalReturns/onlyLocalReturnLambda.kt

@@ -0,0 +1,44 @@
+// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE -NOTHING_TO_INLINE
+import kotlin.InlineOption.*
+
+fun <R> fun1(p: () -> R) {
+    inlineFun {
+        p()
+    }
+}
+
+fun <R> fun1ValueArgument(p: () -> R) {
+    inlineFun ({
+                   p()
+               })
+}
+
+fun <R> fun3(p: () -> R) {
+    inlineFun {
+        <!RETURN_NOT_ALLOWED!>return<!>;
+    }
+}
+
+fun <R> fun3ValueArgument(p: () -> R) {
+    inlineFun ({
+                   <!RETURN_NOT_ALLOWED!>return<!>;
+               })
+}
+
+
+fun <R> fun4(p: () -> R) {
+    inlineFun @lambda {(): R ->
+        return@lambda p();
+    }
+}
+
+fun <R> fun4ValueArgument(p: () -> R) {
+    inlineFun (@lambda {(): R ->
+        return@lambda p();
+    })
+}
+
+
+inline fun <R> inlineFun(inlineOptions(ONLY_LOCAL_RETURN) p: () -> R) {
+    p()
+}

compiler/testData/diagnostics/tests/inline/nonLocalReturns/onlyLocalReturnLambdaBinaryExpr.kt

@@ -0,0 +1,26 @@
+// !DIAGNOSTICS: -UNUSED_EXPRESSION -UNUSED_PARAMETER -UNUSED_VARIABLE -NOTHING_TO_INLINE
+import kotlin.InlineOption.*
+
+class Z {
+    inline fun <R> inlineFun(inlineOptions(ONLY_LOCAL_RETURN) p: () -> R) {
+        p()
+    }
+}
+
+fun <R> fun1(p: () -> R) {
+    Z() inlineFun {
+        p()
+    }
+}
+
+fun <R> fun3(p: () -> R) {
+    Z() inlineFun {
+        <!RETURN_NOT_ALLOWED!>return<!>;
+    }
+}
+
+fun <R> fun4(p: () -> R) {
+    Z() inlineFun @lambda {(): R ->
+        return@lambda p();
+    }
+}

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

@@ -4679,6 +4679,11 @@ public class JetDiagnosticsTestGenerated extends AbstractJetDiagnosticsTest {
                     doTest("compiler/testData/diagnostics/tests/inline/nonLocalReturns/labeledReturn.kt");
                 }
                 
+                @TestMetadata("lambdaWithGlobalReturnsInsideOnlyLocalOne.kt")
+                public void testLambdaWithGlobalReturnsInsideOnlyLocalOne() throws Exception {
+                    doTest("compiler/testData/diagnostics/tests/inline/nonLocalReturns/lambdaWithGlobalReturnsInsideOnlyLocalOne.kt");
+                }
+                
                 @TestMetadata("localFun.kt")
                 public void testLocalFun() throws Exception {
                     doTest("compiler/testData/diagnostics/tests/inline/nonLocalReturns/localFun.kt");
@@ -4699,6 +4704,16 @@ public class JetDiagnosticsTestGenerated extends AbstractJetDiagnosticsTest {
                     doTest("compiler/testData/diagnostics/tests/inline/nonLocalReturns/nonInlinedClass.kt");
                 }
                 
+                @TestMetadata("onlyLocalReturnLambda.kt")
+                public void testOnlyLocalReturnLambda() throws Exception {
+                    doTest("compiler/testData/diagnostics/tests/inline/nonLocalReturns/onlyLocalReturnLambda.kt");
+                }
+                
+                @TestMetadata("onlyLocalReturnLambdaBinaryExpr.kt")
+                public void testOnlyLocalReturnLambdaBinaryExpr() throws Exception {
+                    doTest("compiler/testData/diagnostics/tests/inline/nonLocalReturns/onlyLocalReturnLambdaBinaryExpr.kt");
+                }
+                
                 @TestMetadata("propertyAccessorsAndConstructor.kt")
                 public void testPropertyAccessorsAndConstructor() throws Exception {
                     doTest("compiler/testData/diagnostics/tests/inline/nonLocalReturns/propertyAccessorsAndConstructor.kt");