Replace CHECKCAST kotlin/Unit with ARETURN for tail call optimization

Sometimes instead of {POP, GETSTATIC Unit.INSTANCE, ARETURN} sequence
the codegen emits {CHECKCAST Unit, ARETURN} sequence, which breaks tail
call optimization. By replacing CHECKCAST with ARETURN we eliminate
this issue.

 #KT-19790: Fixed

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

@@ -1,22 +1,11 @@
 /*
- * Copyright 2010-2017 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the 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.
  */
 
 package org.jetbrains.kotlin.codegen.coroutines
 
-import org.jetbrains.kotlin.codegen.inline.isReturnsUnitMarker
+import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.optimization.boxing.isUnitInstance
 import org.jetbrains.kotlin.codegen.optimization.common.ControlFlowGraph
 import org.jetbrains.kotlin.codegen.optimization.common.asSequence
@@ -34,37 +23,60 @@ import org.jetbrains.org.objectweb.asm.tree.analysis.SourceInterpreter
  * Replace POP with ARETURN iff
  * 1) It is immediately followed by { GETSTATIC Unit.INSTANCE, ARETURN } sequences
  * 2) It is poping Unit
+ *
+ * Replace CHECKCAST Unit whit ARETURN iff
+ * It is successed by ARETURN and it casts Unit
  */
 object ReturnUnitMethodTransformer : MethodTransformer() {
     override fun transform(internalClassName: String, methodNode: MethodNode) {
         val unitMarks = findReturnsUnitMarks(methodNode)
         if (unitMarks.isEmpty()) return
 
-        val units = findReturnUnitSequences(methodNode)
-        if (units.isEmpty()) {
-            cleanUpReturnsUnitMarkers(methodNode, unitMarks)
-            return
-        }
+        replaceCheckcastUnitWithAreturn(methodNode, internalClassName)
+        replacePopWithAreturn(methodNode, internalClassName)
 
-        val pops = methodNode.instructions.asSequence().filter { it.opcode == Opcodes.POP }.toList()
-        val popSuccessors = findSuccessors(methodNode, pops)
-        val sourceInsns = findSourceInstructions(internalClassName, methodNode, pops, ignoreCopy = true)
-        val safePops = filterOutUnsafes(popSuccessors, units, sourceInsns)
-
-        // Replace POP with ARETURN for tail call optimization
-        safePops.forEach { methodNode.instructions.set(it, InsnNode(Opcodes.ARETURN)) }
         cleanUpReturnsUnitMarkers(methodNode, unitMarks)
     }
 
-    // Return list of POPs, which can be safely replaced by ARETURNs
+    private fun replacePopWithAreturn(
+        methodNode: MethodNode,
+        internalClassName: String
+    ) {
+        val units = findReturnUnitSequences(methodNode)
+        if (units.isEmpty()) return
+
+        replaceSafeInsnsWithUnit(methodNode, internalClassName, units) { it.opcode == Opcodes.POP }
+    }
+
+    private fun replaceCheckcastUnitWithAreturn(methodNode: MethodNode, internalClassName: String) {
+        val areturns = methodNode.instructions.asSequence().filter { it.opcode == Opcodes.ARETURN }.toList()
+
+        replaceSafeInsnsWithUnit(methodNode, internalClassName, areturns) { it.opcode == Opcodes.CHECKCAST }
+    }
+
+    // Find all instructions, which can be safely replaced with ARETURN and replace them with ARETURN for tail-call optimization
+    private fun replaceSafeInsnsWithUnit(
+        methodNode: MethodNode,
+        internalClassName: String,
+        safeSuccessors: Collection<AbstractInsnNode>,
+        predicate: (AbstractInsnNode) -> Boolean
+    ) {
+        val insns = methodNode.instructions.asSequence().filter(predicate).toList()
+        val successors = findSuccessors(methodNode, insns)
+        val sourceInsns = findSourceInstructions(internalClassName, methodNode, insns, ignoreCopy = true)
+        val safeInsns = filterOutUnsafes(successors, safeSuccessors, sourceInsns)
+        safeInsns.forEach { methodNode.instructions.set(it, InsnNode(Opcodes.ARETURN)) }
+    }
+
+    // Return list of instructions, which can be safely replaced by ARETURNs
     private fun filterOutUnsafes(
-        popSuccessors: Map<AbstractInsnNode, Collection<AbstractInsnNode>>,
-        units: Collection<AbstractInsnNode>,
-        sourceInsns: Map<AbstractInsnNode, Collection<AbstractInsnNode>>
+        successors: Map<AbstractInsnNode, Collection<AbstractInsnNode>>,
+        safeSuccessors: Collection<AbstractInsnNode>,
+        sources: Map<AbstractInsnNode, Collection<AbstractInsnNode>>
     ): Collection<AbstractInsnNode> {
-        return popSuccessors.filter { (pop, successors) ->
-            successors.all { it in units } &&
-                    sourceInsns[pop].sure { "Sources of $pop cannot be null" }.all(::isSuspendingCallReturningUnit)
+        return successors.filter { (insn, successors) ->
+            successors.all { it in safeSuccessors } &&
+                    sources[insn].sure { "Sources of $insn cannot be null" }.all(::isSuspendingCallReturningUnit)
         }.keys
     }
 
@@ -72,7 +84,7 @@ object ReturnUnitMethodTransformer : MethodTransformer() {
     // Return map {insn => list of found instructions}
     private fun findSuccessors(
         methodNode: MethodNode,
-        insns: List<AbstractInsnNode>
+        insns: Collection<AbstractInsnNode>
     ): Map<AbstractInsnNode, Collection<AbstractInsnNode>> {
         val cfg = ControlFlowGraph.build(methodNode)
 

compiler/testData/codegen/box/coroutines/tailCallOptimizations/whenUnit.kt

@@ -0,0 +1,33 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// CHECK_BYTECODE_LISTING
+import helpers.*
+import kotlin.coroutines.experimental.*
+
+sealed class X {
+    class A : X()
+    class B : X()
+}
+
+var log = ""
+
+suspend fun process(a: X.A) { log = "${log}from A;" }
+suspend fun process(b: X.B) { log = "${log}from B;" }
+
+suspend fun process(x: X) = when (x) {
+    is X.A -> process(x)
+    is X.B -> process(x)
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun box(): String {
+    builder {
+        process(X.A())
+        process(X.B())
+    }
+    if (log != "from A;from B;") return log
+    return "OK"
+}

compiler/testData/codegen/box/coroutines/tailCallOptimizations/whenUnit.txt

@@ -0,0 +1,43 @@
+@kotlin.Metadata
+final class WhenUnitKt$box$1 {
+    inner class WhenUnitKt$box$1
+    method <init>(p0: kotlin.coroutines.experimental.Continuation): void
+    public final @org.jetbrains.annotations.NotNull method create(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.experimental.Continuation): kotlin.coroutines.experimental.Continuation
+    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
+    public final @org.jetbrains.annotations.Nullable method invoke(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.experimental.Continuation): java.lang.Object
+    public synthetic method invoke(p0: java.lang.Object): java.lang.Object
+}
+
+@kotlin.Metadata
+public final class WhenUnitKt {
+    private static @org.jetbrains.annotations.NotNull field log: java.lang.String
+    inner class WhenUnitKt$box$1
+    static method <clinit>(): void
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void
+    public final static @org.jetbrains.annotations.NotNull method getLog(): java.lang.String
+    public final static @org.jetbrains.annotations.Nullable method process(@org.jetbrains.annotations.NotNull p0: X$A, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.experimental.Continuation): java.lang.Object
+    public final static @org.jetbrains.annotations.Nullable method process(@org.jetbrains.annotations.NotNull p0: X$B, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.experimental.Continuation): java.lang.Object
+    public final static @org.jetbrains.annotations.Nullable method process(@org.jetbrains.annotations.NotNull p0: X, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.experimental.Continuation): java.lang.Object
+    public final static method setLog(@org.jetbrains.annotations.NotNull p0: java.lang.String): void
+}
+
+@kotlin.Metadata
+public final class X$A {
+    inner class X$A
+    public method <init>(): void
+}
+
+@kotlin.Metadata
+public final class X$B {
+    inner class X$B
+    public method <init>(): void
+}
+
+@kotlin.Metadata
+public abstract class X {
+    inner class X$A
+    inner class X$B
+    private method <init>(): void
+    public synthetic method <init>(p0: kotlin.jvm.internal.DefaultConstructorMarker): void
+}

compiler/testData/codegen/bytecodeListing/tailcall/whenUnit.kt

@@ -0,0 +1,12 @@
+sealed class X {
+    class A : X()
+    class B : X()
+}
+
+suspend fun process(a: X.A) {}
+suspend fun process(b: X.B) {}
+
+suspend fun process(x: X) = when (x) {
+    is X.A -> process(x)
+    is X.B -> process(x)
+}

compiler/testData/codegen/bytecodeListing/tailcall/whenUnit.txt

@@ -0,0 +1,26 @@
+@kotlin.Metadata
+public final class WhenUnitKt {
+    public final static @org.jetbrains.annotations.Nullable method process(@org.jetbrains.annotations.NotNull p0: X$A, @org.jetbrains.annotations.Nullable p1: java.lang.Object): java.lang.Object
+    public final static @org.jetbrains.annotations.Nullable method process(@org.jetbrains.annotations.NotNull p0: X$B, @org.jetbrains.annotations.Nullable p1: java.lang.Object): java.lang.Object
+    public final static @org.jetbrains.annotations.Nullable method process(@org.jetbrains.annotations.NotNull p0: X, @org.jetbrains.annotations.Nullable p1: java.lang.Object): java.lang.Object
+}
+
+@kotlin.Metadata
+public final class X$A {
+    inner class X$A
+    public method <init>(): void
+}
+
+@kotlin.Metadata
+public final class X$B {
+    inner class X$B
+    public method <init>(): void
+}
+
+@kotlin.Metadata
+public abstract class X {
+    inner class X$A
+    inner class X$B
+    private method <init>(): void
+    public synthetic method <init>(p0: kotlin.jvm.internal.DefaultConstructorMarker): void
+}

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

@@ -6836,6 +6836,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unreachable.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("whenUnit.kt")
+            public void testWhenUnit() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tailCallOptimizations/whenUnit.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/tailOperations")

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

@@ -6836,6 +6836,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unreachable.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("whenUnit.kt")
+            public void testWhenUnit() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tailCallOptimizations/whenUnit.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/tailOperations")

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

@@ -157,30 +157,6 @@ public class BytecodeListingTestGenerated extends AbstractBytecodeListingTest {
         doTest(fileName);
     }
 
-    @TestMetadata("tailCallIntrinsics.kt")
-    public void testTailCallIntrinsics() throws Exception {
-        String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/tailCallIntrinsics.kt");
-        doTest(fileName);
-    }
-
-    @TestMetadata("tailSuspendUnitFun.kt")
-    public void testTailSuspendUnitFun() throws Exception {
-        String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/tailSuspendUnitFun.kt");
-        doTest(fileName);
-    }
-
-    @TestMetadata("tryCatchTailCall.kt")
-    public void testTryCatchTailCall() throws Exception {
-        String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/tryCatchTailCall.kt");
-        doTest(fileName);
-    }
-
-    @TestMetadata("unreachable.kt")
-    public void testUnreachable() throws Exception {
-        String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/unreachable.kt");
-        doTest(fileName);
-    }
-
     @TestMetadata("compiler/testData/codegen/bytecodeListing/annotations")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
@@ -357,4 +333,43 @@ public class BytecodeListingTestGenerated extends AbstractBytecodeListingTest {
             }
         }
     }
+
+    @TestMetadata("compiler/testData/codegen/bytecodeListing/tailcall")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Tailcall extends AbstractBytecodeListingTest {
+        public void testAllFilesPresentInTailcall() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/bytecodeListing/tailcall"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY, true);
+        }
+
+        @TestMetadata("tailCallIntrinsics.kt")
+        public void testTailCallIntrinsics() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/tailcall/tailCallIntrinsics.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("tailSuspendUnitFun.kt")
+        public void testTailSuspendUnitFun() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/tailcall/tailSuspendUnitFun.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("tryCatchTailCall.kt")
+        public void testTryCatchTailCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/tailcall/tryCatchTailCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("unreachable.kt")
+        public void testUnreachable() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/tailcall/unreachable.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("whenUnit.kt")
+        public void testWhenUnit() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeListing/tailcall/whenUnit.kt");
+            doTest(fileName);
+        }
+    }
 }

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

@@ -6836,6 +6836,12 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unreachable.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("whenUnit.kt")
+            public void testWhenUnit() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tailCallOptimizations/whenUnit.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/tailOperations")

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -8137,6 +8137,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tailCallOptimizations/unreachable.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("whenUnit.kt")
+            public void testWhenUnit() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tailCallOptimizations/whenUnit.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/tailOperations")