Fix OOM error in ReturnUnitMethodTransformer

#KT-22345: Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/ReturnUnitMethodTransformer.kt

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.codegen.coroutines
 
-import org.jetbrains.kotlin.backend.common.pop
 import org.jetbrains.kotlin.codegen.inline.isReturnsUnitMarker
 import org.jetbrains.kotlin.codegen.optimization.boxing.isUnitInstance
 import org.jetbrains.kotlin.codegen.optimization.common.ControlFlowGraph
@@ -75,31 +74,33 @@ object ReturnUnitMethodTransformer : MethodTransformer() {
         insns: List<AbstractInsnNode>
     ): Map<AbstractInsnNode, Collection<AbstractInsnNode>> {
         val cfg = ControlFlowGraph.build(methodNode)
-        return insns.keysToMap { findSuccessors(cfg, it, methodNode) }
+
+        return insns.keysToMap { findSuccessorsDFS(it, cfg, methodNode) }
     }
 
-    // Find all meaningful successors of [insn]
-    private fun findSuccessors(cfg: ControlFlowGraph, insn: AbstractInsnNode, methodNode: MethodNode): Collection<AbstractInsnNode> {
-        val stack = cfg.getSuccessorsIndices(insn).mapTo(ArrayList()) { methodNode.instructions[it] }
-        val successors = arrayListOf<AbstractInsnNode>()
-        while (stack.isNotEmpty()) {
-            val current = stack.pop()
-            if (current in successors || isReturnsUnitMarker(current)) continue
-            if (!current.isMeaningful || current is JumpInsnNode || current.opcode == Opcodes.NOP) {
-                cfg.getSuccessorsIndices(current).mapTo(stack) { methodNode.instructions[it] }
-                continue
-            }
-            // There can be multiple chains of { UnitInstance, POP } after inlining. Ignore them
-            if (current.isUnitInstance()) {
-                val newSuccessors = findSuccessors(cfg, current, methodNode)
-                if (newSuccessors.all { it.opcode == Opcodes.POP }) {
-                    newSuccessors.flatMapTo(stack) { findSuccessors(cfg, it, methodNode) }
-                    continue
+    // Find all meaningful successors of insn
+    private fun findSuccessorsDFS(insn: AbstractInsnNode, cfg: ControlFlowGraph, methodNode: MethodNode): Collection<AbstractInsnNode> {
+        val visited = hashSetOf<AbstractInsnNode>()
+
+        fun dfs(current: AbstractInsnNode): Collection<AbstractInsnNode> {
+            if (!visited.add(current)) return emptySet()
+
+            return cfg.getSuccessorsIndices(current).flatMap {
+                val succ = methodNode.instructions[it]
+                when {
+                    !succ.isMeaningful || succ is JumpInsnNode || succ.opcode == Opcodes.NOP -> dfs(succ)
+                    succ.isUnitInstance() -> {
+                        // There can be multiple chains of { UnitInstance, POP } after inlining. Ignore them
+                        val newSuccessors = dfs(succ)
+                        if (newSuccessors.all { it.opcode == Opcodes.POP }) newSuccessors.flatMap { dfs(it) }
+                        else setOf(succ)
+                    }
+                    else -> setOf(succ)
                 }
             }
-            successors.add(current)
         }
-        return successors
+
+        return dfs(insn)
     }
 
     private fun isSuspendingCallReturningUnit(node: AbstractInsnNode): Boolean =

compiler/testData/codegen/bytecodeListing/oomInReturnUnit.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+import helpers.*
+import kotlin.coroutines.experimental.*
+import kotlin.coroutines.experimental.intrinsics.*
+
+suspend fun some() {}
+
+suspend fun test() {
+    try {
+        some()
+    } finally {
+        some()
+    }
+}

compiler/testData/codegen/bytecodeListing/oomInReturnUnit.txt

@@ -0,0 +1,71 @@
+@kotlin.Metadata
+final class OomInReturnUnitKt$test$1 {
+    field L$0: java.lang.Object
+    synthetic field data: java.lang.Object
+    synthetic field exception: java.lang.Throwable
+    inner class OomInReturnUnitKt$test$1
+    method <init>(p0: kotlin.coroutines.experimental.Continuation): void
+    public final @org.jetbrains.annotations.Nullable method doResume(@org.jetbrains.annotations.Nullable p0: java.lang.Object, @org.jetbrains.annotations.Nullable p1: java.lang.Throwable): java.lang.Object
+    synthetic final method getLabel(): int
+    synthetic final method setLabel(p0: int): void
+}
+
+@kotlin.Metadata
+public final class OomInReturnUnitKt {
+    inner class OomInReturnUnitKt$test$1
+    public final static @org.jetbrains.annotations.Nullable method some(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.experimental.Continuation): java.lang.Object
+    public final static @org.jetbrains.annotations.Nullable method test(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.experimental.Continuation): java.lang.Object
+}
+
+@kotlin.Metadata
+public final class helpers/CoroutineUtilKt$handleExceptionContinuation$1 {
+    synthetic final field $x: kotlin.jvm.functions.Function1
+    private final @org.jetbrains.annotations.NotNull field context: kotlin.coroutines.experimental.EmptyCoroutineContext
+    inner class helpers/CoroutineUtilKt$handleExceptionContinuation$1
+    method <init>(p0: kotlin.jvm.functions.Function1): void
+    public synthetic method getContext(): kotlin.coroutines.experimental.CoroutineContext
+    public @org.jetbrains.annotations.NotNull method getContext(): kotlin.coroutines.experimental.EmptyCoroutineContext
+    public method resume(@org.jetbrains.annotations.Nullable p0: java.lang.Object): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+}
+
+@kotlin.Metadata
+public final class helpers/CoroutineUtilKt$handleResultContinuation$1 {
+    synthetic final field $x: kotlin.jvm.functions.Function1
+    private final @org.jetbrains.annotations.NotNull field context: kotlin.coroutines.experimental.EmptyCoroutineContext
+    inner class helpers/CoroutineUtilKt$handleResultContinuation$1
+    method <init>(p0: kotlin.jvm.functions.Function1): void
+    public synthetic method getContext(): kotlin.coroutines.experimental.CoroutineContext
+    public @org.jetbrains.annotations.NotNull method getContext(): kotlin.coroutines.experimental.EmptyCoroutineContext
+    public method resume(p0: java.lang.Object): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+}
+
+@kotlin.Metadata
+public final class helpers/CoroutineUtilKt {
+    inner class helpers/CoroutineUtilKt$handleExceptionContinuation$1
+    inner class helpers/CoroutineUtilKt$handleResultContinuation$1
+    public final static @org.jetbrains.annotations.NotNull method handleExceptionContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.experimental.Continuation
+    public final static @org.jetbrains.annotations.NotNull method handleResultContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.experimental.Continuation
+}
+
+@kotlin.Metadata
+public final class helpers/EmptyContinuation$Companion {
+    inner class helpers/EmptyContinuation$Companion
+    private method <init>(): void
+    public synthetic method <init>(p0: kotlin.jvm.internal.DefaultConstructorMarker): void
+}
+
+@kotlin.Metadata
+public class helpers/EmptyContinuation {
+    public final static field Companion: helpers.EmptyContinuation$Companion
+    private final @org.jetbrains.annotations.NotNull field context: kotlin.coroutines.experimental.CoroutineContext
+    inner class helpers/EmptyContinuation$Companion
+    static method <clinit>(): void
+    public method <init>(): void
+    public method <init>(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.experimental.CoroutineContext): void
+    public synthetic method <init>(p0: kotlin.coroutines.experimental.CoroutineContext, p1: int, p2: kotlin.jvm.internal.DefaultConstructorMarker): void
+    public @org.jetbrains.annotations.NotNull method getContext(): kotlin.coroutines.experimental.CoroutineContext
+    public method resume(@org.jetbrains.annotations.Nullable p0: java.lang.Object): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+}

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

@@ -139,6 +139,12 @@ public class BytecodeListingTestGenerated extends AbstractBytecodeListingTest {
         doTest(fileName);
     }
 
+    @TestMetadata("oomInReturnUnit.kt")
+    public void testOomInReturnUnit() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/oomInReturnUnit.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("samAdapterAndInlinedOne.kt")
     public void testSamAdapterAndInlinedOne() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/samAdapterAndInlinedOne.kt");