Support recursive local suspend functions

Previously they were trying to call constructor (incorrectly) for
recursive calls. This is redundant, since this.invoke creates one
automatically.

In addition, support callable references to recursive local suspend
functions.

 #KT-24780 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/StackValue.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
  * that can be found in the license/LICENSE.txt file.
  */
 
@@ -15,6 +15,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
 import org.jetbrains.kotlin.builtins.PrimitiveType;
+import org.jetbrains.kotlin.codegen.binding.CodegenBinding;
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineCodegenUtilKt;
 import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethods;
 import org.jetbrains.kotlin.codegen.pseudoInsns.PseudoInsnsKt;
@@ -49,6 +50,8 @@ import java.util.List;
 import java.util.function.Consumer;
 
 import static org.jetbrains.kotlin.codegen.AsmUtil.*;
+import static org.jetbrains.kotlin.codegen.binding.CodegenBinding.CLASS_FOR_CALLABLE;
+import static org.jetbrains.kotlin.codegen.binding.CodegenBinding.RECURSIVE_SUSPEND_CALLABLE_REFERENCE;
 import static org.jetbrains.kotlin.resolve.jvm.AsmTypes.*;
 import static org.jetbrains.org.objectweb.asm.Opcodes.*;
 
@@ -850,12 +853,7 @@ public abstract class StackValue {
                     SimpleFunctionDescriptor initial =
                             CoroutineCodegenUtilKt.unwrapInitialDescriptorForSuspendFunction((SimpleFunctionDescriptor) descriptor);
                     if (initial != null && initial.isSuspend()) {
-                        StackValue value = codegen.findLocalOrCapturedValue(initial.getOriginal());
-                        assert value != null : "Local suspend fun should be found in locals or in captured params: " +
-                                               descriptor +
-                                               " initial local suspend fun: " +
-                                               initial;
-                        return value;
+                        return putLocalSuspendFunctionOnStack(codegen, initial.getOriginal());
                     }
                 }
                 StackValue value = codegen.findLocalOrCapturedValue(descriptor.getOriginal());
@@ -873,6 +871,48 @@ public abstract class StackValue {
         return none();
     }
 
+    private static StackValue putLocalSuspendFunctionOnStack(
+            @NotNull ExpressionCodegen codegen,
+            SimpleFunctionDescriptor callee
+    ) {
+        // There can be three types of suspend local function calls:
+        // 1) normal call: we first define it as a closure and then call it
+        // 2) call using callable reference: in this case it is not local, but rather captured value
+        // 3) recursive call: we are in the middle of defining it, but, thankfully, we can simply call `this.invoke` to
+        // create new coroutine
+
+        // First, check whether this is a normal call
+        int index = codegen.lookupLocalIndex(callee);
+        if (index >= 0) {
+            // This is a normal local call
+            return local(index, OBJECT_TYPE);
+        }
+
+        // Then check for call inside a callable reference
+        BindingContext bindingContext = codegen.getBindingContext();
+        Type calleeType = CodegenBinding.asmTypeForAnonymousClass(bindingContext, callee);
+        if (codegen.context.hasThisDescriptor()) {
+            ClassDescriptor thisDescriptor = codegen.context.getThisDescriptor();
+            if (thisDescriptor instanceof SyntheticClassDescriptorForLambda &&
+                ((SyntheticClassDescriptorForLambda) thisDescriptor).isCallableReference()) {
+                // Call is inside a callable reference
+                // if it is call to recursive local, just return this$0
+                Boolean isRecursive = bindingContext.get(RECURSIVE_SUSPEND_CALLABLE_REFERENCE, thisDescriptor);
+                if (isRecursive != null && isRecursive) {
+                    ClassDescriptor classDescriptor =
+                            bindingContext.get(CLASS_FOR_CALLABLE, callee);
+                    assert classDescriptor != null : "No CLASS_FOR_CALLABLE" + callee;
+                    return thisOrOuter(codegen, classDescriptor, false, false);
+                }
+                // Otherwise, just call constructor of the closure
+                return codegen.findCapturedValue(callee);
+            }
+        }
+        // Recursive suspend local function, just call invoke on this, it will create new coroutine automatically
+        codegen.v.visitVarInsn(ALOAD, 0);
+        return onStack(calleeType);
+    }
+
     private static StackValue platformStaticCallIfPresent(@NotNull StackValue resultReceiver, @NotNull CallableDescriptor descriptor) {
         if (CodegenUtilKt.isJvmStaticInObjectOrClassOrInterface(descriptor)) {
             if (resultReceiver.canHaveSideEffects()) {

compiler/backend/src/org/jetbrains/kotlin/codegen/binding/CodegenAnnotatingVisitor.java

@@ -399,6 +399,14 @@ class CodegenAnnotatingVisitor extends KtVisitorVoid {
             SimpleFunctionDescriptor functionDescriptor = (SimpleFunctionDescriptor) callableDescriptor;
             if (functionDescriptor.isSuspend()) {
                 createAndRecordSuspendFunctionView(closure, functionDescriptor, false);
+                boolean isRecursive = false;
+                for (FunctionDescriptor descriptor: functionsStack) {
+                    if (descriptor == target) {
+                        isRecursive = true;
+                        break;
+                    }
+                }
+                bindingTrace.record(RECURSIVE_SUSPEND_CALLABLE_REFERENCE, classDescriptor, isRecursive);
             }
         }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/binding/CodegenBinding.java

@@ -58,6 +58,9 @@ public class CodegenBinding {
     public static final WritableSlice<FunctionDescriptor, Boolean> CAPTURES_CROSSINLINE_SUSPEND_LAMBDA =
             Slices.createSimpleSlice();
 
+    public static final WritableSlice<ClassDescriptor, Boolean> RECURSIVE_SUSPEND_CALLABLE_REFERENCE =
+            Slices.createSimpleSlice();
+
     public static final WritableSlice<ValueParameterDescriptor, ValueParameterDescriptor> PARAMETER_SYNONYM =
             Slices.createSimpleSlice();
 

compiler/testData/codegen/box/coroutines/localFunctions/named/rec.kt

@@ -0,0 +1,90 @@
+// LANGUAGE_VERSION: 1.3
+// WITH_RUNTIME
+// WITH_COROUTINES
+// TARGET_BACKEND: JVM
+// IGNORE_BACKEND: JVM_IR
+
+import helpers.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(handleResultContinuation {
+        proceed = null
+    })
+}
+
+suspend fun foo(until: Int): String {
+    val o = "O"
+    val k = "K"
+    val dot = "."
+    suspend fun bar(x: Int): String =
+        if (x == until) dot else if (x < until) o + bar(x * 2) else k + bar(x - 1)
+    return bar(1)
+}
+
+var proceed: (() -> Unit)? = null
+
+suspend fun suspendHere(value: String): String = suspendCoroutineUninterceptedOrReturn { x ->
+    proceed = {
+        x.resume(value)
+    }
+    COROUTINE_SUSPENDED
+}
+
+suspend fun foo2(until: Int): String {
+    val o = "O"
+    val k = "K"
+    val dot = "."
+    suspend fun bar(x: Int): String =
+        if (x == until) dot else if (x < until) suspendHere(o + bar(x * 2)) else suspendHere(k + bar(x - 1))
+    return bar(1)
+}
+
+suspend fun fooCallableReference(until: Int): String {
+    val o = "O"
+    val k = "K"
+    val dot = "."
+    suspend fun bar(x: Int): String =
+        if (x == until) dot else if (x < until) o + (::bar)(x * 2) else k + (::bar)(x - 1)
+    return bar(1)
+}
+
+suspend fun fooCallableReferenceIndirectRecursion(until: Int): String {
+    val o = "O"
+    val k = "K"
+    val dot = "."
+    suspend fun bar(x: Int): String {
+        suspend fun innerO() = o + (::bar)(x * 2)
+        suspend fun innerK() = k + (::bar)(x - 1)
+        return if (x == until) dot else if (x < until) innerO() else innerK()
+    }
+    return bar(1)
+}
+
+fun box(): String {
+    var res = "FAIL 1"
+    builder {
+        res = foo(10)
+    }
+    if (res != "OOOOKKKKKK.") return "FAIL 1: $res"
+    res = "FAIL 2"
+    builder {
+        res = foo2(10)
+    }
+    while (proceed != null) {
+        proceed!!()
+    }
+    if (res != "OOOOKKKKKK.") return "FAIL 2: $res"
+    res = "FAIL 3"
+    builder {
+        res = fooCallableReference(10)
+    }
+    if (res != "OOOOKKKKKK.") return "FAIL 3: $res"
+    res = "FAIL 4"
+    builder {
+        res = fooCallableReferenceIndirectRecursion(10)
+    }
+    if (res != "OOOOKKKKKK.") return "FAIL 4: $res"
+    return "OK"
+}

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -7311,6 +7311,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                     runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/localFunctions/named/nestedLocals.kt", "kotlin.coroutines");
                 }
 
+                @TestMetadata("rec.kt")
+                public void testRec() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/localFunctions/named/rec.kt");
+                }
+
                 @TestMetadata("simpleSuspensionPoint.kt")
                 public void testSimpleSuspensionPoint_1_2() throws Exception {
                     runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/localFunctions/named/simpleSuspensionPoint.kt", "kotlin.coroutines.experimental");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -7311,6 +7311,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                     runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/localFunctions/named/nestedLocals.kt", "kotlin.coroutines");
                 }
 
+                @TestMetadata("rec.kt")
+                public void testRec() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/localFunctions/named/rec.kt");
+                }
+
                 @TestMetadata("simpleSuspensionPoint.kt")
                 public void testSimpleSuspensionPoint_1_2() throws Exception {
                     runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/localFunctions/named/simpleSuspensionPoint.kt", "kotlin.coroutines.experimental");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -7311,6 +7311,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                     runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/localFunctions/named/nestedLocals.kt", "kotlin.coroutines");
                 }
 
+                @TestMetadata("rec.kt")
+                public void testRec() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/localFunctions/named/rec.kt");
+                }
+
                 @TestMetadata("simpleSuspensionPoint.kt")
                 public void testSimpleSuspensionPoint_1_2() throws Exception {
                     runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/localFunctions/named/simpleSuspensionPoint.kt", "kotlin.coroutines.experimental");