More wise lambda search during inlining

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InstructionsAndFrames.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2016 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.
+ */
+
+package org.jetbrains.kotlin.codegen.inline
+
+import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
+import org.jetbrains.org.objectweb.asm.tree.InsnList
+import org.jetbrains.org.objectweb.asm.tree.analysis.Frame
+import org.jetbrains.org.objectweb.asm.tree.analysis.SourceValue
+
+class InstructionsAndFrames(val frames: Array<Frame<SourceValue>>, ins: InsnList) {
+
+    val indexes = (0..ins.size() - 1).associate { Pair(ins.get(it), it) }
+
+    operator fun get(node: AbstractInsnNode): Frame<SourceValue>? = frames[indexes[node]!!]
+
+}

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/MethodInliner.java

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.codegen.inline;
 
 import com.google.common.collect.Lists;
 import com.intellij.util.ArrayUtil;
+import com.intellij.util.containers.SmartHashSet;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.codegen.ClosureCodegen;
@@ -427,7 +428,8 @@ public class MethodInliner {
         catch (AnalyzerException e) {
             throw wrapException(e, node, "couldn't inline method call");
         }
-
+        Set<AbstractInsnNode> toDelete = new SmartHashSet<AbstractInsnNode>();
+        InstructionsAndFrames instructionsAndFrames = new InstructionsAndFrames(sources, node.instructions);
         AbstractInsnNode cur = node.instructions.getFirst();
         int index = 0;
 
@@ -468,7 +470,7 @@ public class MethodInliner {
                             if (insnNode.getOpcode() == Opcodes.SWAP) {
                                 processingInstruction = InlineCodegenUtil.getPrevMeaningful(insnNode);
                             }
-                            lambdaInfo = getLambdaIfExists(processingInstruction);
+                            lambdaInfo = getLambdaIfExistsAndMarkInstructions(processingInstruction, frame, instructionsAndFrames, toDelete);
                             if (lambdaInfo != null) {
                                 //remove inlinable access
                                 assert processingInstruction != null;
@@ -486,7 +488,7 @@ public class MethodInliner {
                             SourceValue sourceValue = frame.getStack(firstParameterIndex + i);
                             if (sourceValue.insns.size() == 1) {
                                 AbstractInsnNode insnNode = sourceValue.insns.iterator().next();
-                                LambdaInfo lambdaInfo = getLambdaIfExists(insnNode);
+                                LambdaInfo lambdaInfo = getLambdaIfExistsAndMarkInstructions(insnNode, frame, instructionsAndFrames, toDelete);
                                 if (lambdaInfo != null) {
                                     lambdaMapping.put(offset, lambdaInfo);
                                     node.instructions.remove(insnNode);
@@ -533,6 +535,10 @@ public class MethodInliner {
             }
         }
 
+        for (AbstractInsnNode insnNode : toDelete) {
+            node.instructions.remove(insnNode);
+        }
+
         //clean dead try/catch blocks
         List<TryCatchBlockNode> blocks = node.tryCatchBlocks;
         for (Iterator<TryCatchBlockNode> iterator = blocks.iterator(); iterator.hasNext(); ) {
@@ -575,7 +581,41 @@ public class MethodInliner {
     }
 
     @Nullable
-    public LambdaInfo getLambdaIfExists(@Nullable AbstractInsnNode insnNode) {
+    private LambdaInfo getLambdaIfExistsAndMarkInstructions(
+            @Nullable AbstractInsnNode insnNode,
+            @NotNull Frame<SourceValue> localFrame,
+            @NotNull InstructionsAndFrames insAndFrames,
+            @NotNull Set<AbstractInsnNode> toDelete
+    ) {
+        LambdaInfo lambdaInfo = getLambdaIfExists(insnNode);
+
+        if (lambdaInfo == null && insnNode instanceof VarInsnNode && insnNode.getOpcode() == Opcodes.ALOAD) {
+            int varIndex = ((VarInsnNode) insnNode).var;
+            SourceValue local = localFrame.getLocal(varIndex);
+            if (local.insns.size() == 1) {
+                AbstractInsnNode storeIns = local.insns.iterator().next();
+                if (storeIns instanceof VarInsnNode && storeIns.getOpcode() == Opcodes.ASTORE) {
+                    Frame<SourceValue> frame = insAndFrames.get(storeIns);
+                    if (frame != null) {
+                        SourceValue topOfStack = frame.getStack(frame.getStackSize() - 1);
+                        if(topOfStack.insns.size() == 1) {
+                            AbstractInsnNode lambdaAload = topOfStack.insns.iterator().next();
+                            lambdaInfo = getLambdaIfExistsAndMarkInstructions(lambdaAload, frame, insAndFrames, toDelete);
+                            if (lambdaInfo != null) {
+                                toDelete.add(storeIns);
+                                toDelete.add(lambdaAload);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        return lambdaInfo;
+    }
+
+    @Nullable
+    private LambdaInfo getLambdaIfExists(@Nullable AbstractInsnNode insnNode) {
         if (insnNode == null) {
             return null;
         }

compiler/testData/codegen/boxInline/compexStack/asCheck.1.kt

@@ -0,0 +1,5 @@
+import test.*
+
+fun box(): String {
+    return ContentTypeByExtension.processRecords { ext -> ext }
+}

compiler/testData/codegen/boxInline/compexStack/asCheck.2.kt

@@ -0,0 +1,18 @@
+package test
+
+object ContentTypeByExtension {
+    inline fun processRecords(crossinline operation: (String) -> String) =
+            listOf("O", "K").map {
+                val ext = B(it)
+                operation(ext.toLowerCase())
+            }.joinToString("")
+}
+
+
+
+
+inline fun A.toLowerCase(): String = (this as B).value
+
+open class A
+
+open class B(val value: String) : A()

compiler/testData/codegen/boxInline/compexStack/asCheck2.1.kt

@@ -0,0 +1,5 @@
+import test.*
+
+fun box(): String {
+    return ContentTypeByExtension.processRecords { ext -> ext }
+}

compiler/testData/codegen/boxInline/compexStack/asCheck2.2.kt

@@ -0,0 +1,18 @@
+package test
+
+object ContentTypeByExtension {
+    inline fun processRecords(crossinline operation: (String) -> String) =
+             {
+                val ext = B("OK")
+                operation(ext.toLowerCase())
+            }()
+}
+
+
+
+
+inline fun A.toLowerCase(): String = (this as B).value
+
+open class A
+
+open class B(val value: String) : A()

compiler/testData/codegen/boxInline/compexStack/simple.1.kt

@@ -0,0 +1,5 @@
+import test.*
+
+fun box(): String {
+    return processRecords { ext -> ext + "K" }
+}

compiler/testData/codegen/boxInline/compexStack/simple.2.kt

@@ -0,0 +1,5 @@
+package test
+
+inline fun foo(x: String) = x
+
+inline fun processRecords(block: (String) -> String) = block(foo("O"))

compiler/testData/codegen/boxInline/compexStack/simple2.1.kt

@@ -0,0 +1,5 @@
+import test.*
+
+fun box(): String {
+    return processRecords { "O" + it }
+}

compiler/testData/codegen/boxInline/compexStack/simple2.2.kt

@@ -0,0 +1,11 @@
+package test
+
+inline fun foo(x: String) = x
+
+class A {
+    fun test(s: String) = s
+}
+
+inline fun processRecords(block: (String) -> String): String {
+    return A().test(block(foo("K")))
+}

compiler/testData/codegen/boxInline/compexStack/simple3.1.kt

@@ -0,0 +1,7 @@
+import test.*
+
+fun box(): String {
+    val result = processRecords { "B" + it }
+
+    return if (result == "BOK1") "OK" else "fail: $result"
+}

compiler/testData/codegen/boxInline/compexStack/simple3.2.kt

@@ -0,0 +1,11 @@
+package test
+
+inline fun foo(x: String, y: String) = x + y
+
+class A {
+    fun test(s: String) = s
+}
+
+inline fun processRecords(block: (String) -> String): String {
+    return A().test(block(foo("O", foo("K", "1"))))
+}

compiler/testData/codegen/boxInline/compexStack/simple4.1.kt

@@ -0,0 +1,5 @@
+import test.*
+
+fun box(): String {
+    return processRecords { a, b -> a + b}
+}

compiler/testData/codegen/boxInline/compexStack/simple4.2.kt

@@ -0,0 +1,10 @@
+package test
+
+inline fun foo(x: String) = x
+
+fun test(a: String, s: String) = s
+
+
+inline fun processRecords(block: (String, String) -> String): String {
+    return test("stub", block(foo("O"), foo("K")))
+}

compiler/tests/org/jetbrains/kotlin/codegen/generated/BlackBoxInlineCodegenTestGenerated.java

@@ -548,6 +548,51 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
         }
     }
 
+    @TestMetadata("compiler/testData/codegen/boxInline/compexStack")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class CompexStack extends AbstractBlackBoxInlineCodegenTest {
+        public void testAllFilesPresentInCompexStack() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/compexStack"), Pattern.compile("^(.+)\\.1.kt$"), true);
+        }
+
+        @TestMetadata("asCheck.1.kt")
+        public void testAsCheck() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/asCheck.1.kt");
+            doTestMultiFileWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("asCheck2.1.kt")
+        public void testAsCheck2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/asCheck2.1.kt");
+            doTestMultiFileWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("simple.1.kt")
+        public void testSimple() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple.1.kt");
+            doTestMultiFileWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("simple2.1.kt")
+        public void testSimple2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple2.1.kt");
+            doTestMultiFileWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("simple3.1.kt")
+        public void testSimple3() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple3.1.kt");
+            doTestMultiFileWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("simple4.1.kt")
+        public void testSimple4() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple4.1.kt");
+            doTestMultiFileWithInlineCheck(fileName);
+        }
+    }
+
     @TestMetadata("compiler/testData/codegen/boxInline/complex")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

compiler/tests/org/jetbrains/kotlin/jvm/compiler/CompileKotlinAgainstInlineKotlinTestGenerated.java

@@ -548,6 +548,51 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
         }
     }
 
+    @TestMetadata("compiler/testData/codegen/boxInline/compexStack")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class CompexStack extends AbstractCompileKotlinAgainstInlineKotlinTest {
+        public void testAllFilesPresentInCompexStack() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/compexStack"), Pattern.compile("^(.+)\\.1.kt$"), true);
+        }
+
+        @TestMetadata("asCheck.1.kt")
+        public void testAsCheck() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/asCheck.1.kt");
+            doBoxTestWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("asCheck2.1.kt")
+        public void testAsCheck2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/asCheck2.1.kt");
+            doBoxTestWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("simple.1.kt")
+        public void testSimple() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple.1.kt");
+            doBoxTestWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("simple2.1.kt")
+        public void testSimple2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple2.1.kt");
+            doBoxTestWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("simple3.1.kt")
+        public void testSimple3() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple3.1.kt");
+            doBoxTestWithInlineCheck(fileName);
+        }
+
+        @TestMetadata("simple4.1.kt")
+        public void testSimple4() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple4.1.kt");
+            doBoxTestWithInlineCheck(fileName);
+        }
+    }
+
     @TestMetadata("compiler/testData/codegen/boxInline/complex")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)