Support extension lambda inlining in complex stack cases (nullable receiver)

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

@@ -463,19 +463,11 @@ public class MethodInliner {
                         LambdaInfo lambdaInfo = null;
                         int varIndex = -1;
 
-                        if (sourceValue.insns.size() == 1) {
+                        AbstractInsnNode insnNode = MethodInlinerUtilKt.singleOrNullInsn(sourceValue);
-                            AbstractInsnNode insnNode = sourceValue.insns.iterator().next();
+                        if (insnNode != null) {
-                            AbstractInsnNode processingInstruction = insnNode;
+                            lambdaInfo = MethodInlinerUtilKt.getLambdaIfExistsAndMarkInstructions(
-
+                                    this, insnNode, frame, instructionsAndFrames, toDelete, true
-                            if (insnNode.getOpcode() == Opcodes.SWAP) {
+                            );
-                                processingInstruction = InlineCodegenUtil.getPrevMeaningful(insnNode);
-                            }
-                            lambdaInfo = getLambdaIfExistsAndMarkInstructions(processingInstruction, frame, instructionsAndFrames, toDelete);
-                            if (lambdaInfo != null) {
-                                //remove inlinable access
-                                assert processingInstruction != null;
-                                InlineCodegenUtil.removeInterval(node, processingInstruction, insnNode);
-                            }
                         }
 
                         invokeCalls.add(new InvokeCall(varIndex, lambdaInfo, currentFinallyDeep));
@@ -486,12 +478,13 @@ public class MethodInliner {
                         int offset = 0;
                         for (int i = 0; i < paramCount; i++) {
                             SourceValue sourceValue = frame.getStack(firstParameterIndex + i);
-                            if (sourceValue.insns.size() == 1) {
+                            AbstractInsnNode insnNode = MethodInlinerUtilKt.singleOrNullInsn(sourceValue);
-                                AbstractInsnNode insnNode = sourceValue.insns.iterator().next();
+                            if (insnNode != null) {
-                                LambdaInfo lambdaInfo = getLambdaIfExistsAndMarkInstructions(insnNode, frame, instructionsAndFrames, toDelete);
+                                LambdaInfo lambdaInfo = MethodInlinerUtilKt.getLambdaIfExistsAndMarkInstructions(
+                                        this, insnNode, frame, instructionsAndFrames, toDelete, false
+                                );
                                 if (lambdaInfo != null) {
                                     lambdaMapping.put(offset, lambdaInfo);
-                                    node.instructions.remove(insnNode);
                                 }
                             }
                             offset += i == 0 ? 1 : argTypes[i - 1].getSize();
@@ -581,41 +574,7 @@ public class MethodInliner {
     }
 
     @Nullable
-    private LambdaInfo getLambdaIfExistsAndMarkInstructions(
+    LambdaInfo getLambdaIfExists(@Nullable AbstractInsnNode insnNode) {
-            @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/backend/src/org/jetbrains/kotlin/codegen/inline/MethodInlinerUtil.kt

@@ -0,0 +1,68 @@
+/*
+ * 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.kotlin.codegen.optimization.common.InsnSequence
+import org.jetbrains.org.objectweb.asm.Opcodes
+import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
+import org.jetbrains.org.objectweb.asm.tree.VarInsnNode
+import org.jetbrains.org.objectweb.asm.tree.analysis.Frame
+import org.jetbrains.org.objectweb.asm.tree.analysis.SourceValue
+
+
+fun MethodInliner.getLambdaIfExistsAndMarkInstructions(
+        insnNode: AbstractInsnNode,
+        localFrame: Frame<SourceValue>,
+        insAndFrames: InstructionsAndFrames,
+        toDelete: MutableSet<AbstractInsnNode>,
+        processSwap: Boolean): LambdaInfo? {
+
+    val processingInsnNode = if (processSwap && insnNode.opcode == Opcodes.SWAP) InlineCodegenUtil.getPrevMeaningful(insnNode) else insnNode
+
+    var lambdaInfo = getLambdaIfExists(processingInsnNode)
+
+    if (lambdaInfo == null && processingInsnNode is VarInsnNode && processingInsnNode.opcode == Opcodes.ALOAD) {
+        val varIndex = processingInsnNode.`var`
+        val local = localFrame.getLocal(varIndex)
+        val storeIns = local.singleOrNullInsn()
+        if (storeIns is VarInsnNode && storeIns.getOpcode() == Opcodes.ASTORE) {
+            val frame = insAndFrames[storeIns]
+            if (frame != null) {
+                val topOfStack = frame.getStack(frame.stackSize - 1)
+                val lambdaAload = topOfStack.singleOrNullInsn()
+                if (lambdaAload != null) {
+                    lambdaInfo = getLambdaIfExistsAndMarkInstructions(lambdaAload, frame, insAndFrames, toDelete, processSwap)
+                    if (lambdaInfo != null) {
+                        toDelete.add(storeIns)
+                        toDelete.add(lambdaAload)
+                    }
+                }
+            }
+        }
+    }
+
+    if (lambdaInfo != null) {
+        InsnSequence(processingInsnNode!!, insnNode).forEach { toDelete.add(it) }
+        toDelete.add(insnNode)
+    }
+
+    return lambdaInfo
+}
+
+fun SourceValue.singleOrNullInsn(): AbstractInsnNode? {
+    return insns.singleOrNull()
+}

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

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

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

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

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

@@ -591,6 +591,12 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple4.1.kt");
             doTestMultiFileWithInlineCheck(fileName);
         }
+
+        @TestMetadata("simpleExtension.1.kt")
+        public void testSimpleExtension() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simpleExtension.1.kt");
+            doTestMultiFileWithInlineCheck(fileName);
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/boxInline/complex")

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

@@ -591,6 +591,12 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simple4.1.kt");
             doBoxTestWithInlineCheck(fileName);
         }
+
+        @TestMetadata("simpleExtension.1.kt")
+        public void testSimpleExtension() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxInline/compexStack/simpleExtension.1.kt");
+            doBoxTestWithInlineCheck(fileName);
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/boxInline/complex")