Replace SourceInterpreter with interpreter, which track only

functional arguments.

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

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.codegen.inline.coroutines.CoroutineTransformer
 import org.jetbrains.kotlin.codegen.inline.coroutines.isSuspendLambdaCapturedByOuterObjectOrLambda
 import org.jetbrains.kotlin.codegen.inline.coroutines.markNoinlineLambdaIfSuspend
 import org.jetbrains.kotlin.codegen.inline.coroutines.surroundInvokesWithSuspendMarkersIfNeeded
-import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethods
 import org.jetbrains.kotlin.codegen.optimization.ApiVersionCallsPreprocessingMethodTransformer
 import org.jetbrains.kotlin.codegen.optimization.FixStackWithLabelNormalizationMethodTransformer
 import org.jetbrains.kotlin.codegen.optimization.common.ControlFlowGraph
@@ -501,9 +500,9 @@ class MethodInliner(
 
         replaceContinuationAccessesWithFakeContinuationsIfNeeded(processingNode)
 
-        val sources = analyzeMethodNodeBeforeInline(processingNode)
-
         val toDelete = SmartSet.create<AbstractInsnNode>()
+
+        val sources = analyzeMethodWithFunctionalArgumentInterpreter(this, processingNode, toDelete)
         val instructions = processingNode.instructions
 
         var awaitClassReification = false
@@ -530,9 +529,7 @@ class MethodInliner(
                         val firstParameterIndex = frame.stackSize - paramCount
                         if (isInvokeOnLambda(owner, name) /*&& methodInsnNode.owner.equals(INLINE_RUNTIME)*/) {
                             val sourceValue = frame.getStack(firstParameterIndex)
-                            val functionalArgument =
+                            invokeCalls.add(InvokeCall(sourceValue.functionalArgument, currentFinallyDeep))
-                                getFunctionalArgumentIfExistsAndMarkInstructions(sourceValue, true, instructions, sources, toDelete)
-                            invokeCalls.add(InvokeCall(functionalArgument, currentFinallyDeep))
                         } else if (isSamWrapperConstructorCall(owner, name)) {
                             recordTransformation(SamWrapperTransformationInfo(owner, inliningContext, isAlreadyRegenerated(owner)))
                         } else if (isAnonymousConstructorCall(owner, name)) {
@@ -542,11 +539,9 @@ class MethodInliner(
                             var capturesAnonymousObjectThatMustBeRegenerated = false
                             for (i in 0 until paramCount) {
                                 val sourceValue = frame.getStack(firstParameterIndex + i)
-                                val functionalArgument = getFunctionalArgumentIfExistsAndMarkInstructions(
+                                val functionalArgument = sourceValue.functionalArgument
-                                    sourceValue, false, instructions, sources, toDelete
-                                )
                                 if (functionalArgument != null) {
-                                    functionalArgumentMapping.put(offset, functionalArgument)
+                                    functionalArgumentMapping[offset] = functionalArgument
                                 } else if (i < argTypes.size && isAnonymousClassThatMustBeRegenerated(argTypes[i])) {
                                     capturesAnonymousObjectThatMustBeRegenerated = true
                                 }
@@ -556,7 +551,11 @@ class MethodInliner(
 
                             recordTransformation(
                                 buildConstructorInvocation(
-                                    owner, cur.desc, functionalArgumentMapping, awaitClassReification, capturesAnonymousObjectThatMustBeRegenerated
+                                    owner,
+                                    cur.desc,
+                                    functionalArgumentMapping,
+                                    awaitClassReification,
+                                    capturesAnonymousObjectThatMustBeRegenerated
                                 )
                             )
                             awaitClassReification = false
@@ -604,14 +603,8 @@ class MethodInliner(
                         }
                     }
 
-                    cur.opcode == Opcodes.POP -> getFunctionalArgumentIfExistsAndMarkInstructions(
+                    cur.opcode == Opcodes.POP -> {
-                        frame.top()!!,
+                        if (frame.top().functionalArgument is LambdaInfo) {
-                        true,
-                        instructions,
-                        sources,
-                        toDelete
-                    )?.let {
-                        if (it is LambdaInfo) {
                             toDelete.add(cur)
                         }
                     }
@@ -630,8 +623,7 @@ class MethodInliner(
                             nodeRemapper.originalLambdaInternalName == fieldInsn.owner
                         ) {
                             val stackTransformations = mutableSetOf<AbstractInsnNode>()
-                            val lambdaInfo =
+                            val lambdaInfo = frame.peek(1)?.functionalArgument
-                                getFunctionalArgumentIfExistsAndMarkInstructions(frame.peek(1)!!, false, instructions, sources, stackTransformations)
                             if (lambdaInfo is LambdaInfo && stackTransformations.all { it is VarInsnNode }) {
                                 assert(lambdaInfo.lambdaClassType.internalName == nodeRemapper.originalLambdaInternalName) {
                                     "Wrong bytecode template for contract template: ${lambdaInfo.lambdaClassType.internalName} != ${nodeRemapper.originalLambdaInternalName}"
@@ -1076,9 +1068,7 @@ class MethodInliner(
         private fun removeClosureAssertions(node: MethodNode) {
             val toDelete = arrayListOf<AbstractInsnNode>()
             InsnSequence(node.instructions).filterIsInstance<MethodInsnNode>().forEach { methodInsnNode ->
-                if (methodInsnNode.owner == IntrinsicMethods.INTRINSICS_CLASS_NAME &&
+                if (isParameterNullabilityCheck(methodInsnNode)) {
-                    (methodInsnNode.name == "checkParameterIsNotNull" || methodInsnNode.name == "checkNotNullParameter")
-                ) {
                     val prev = methodInsnNode.previous
                     assert(Opcodes.LDC == prev?.opcode) { "'${methodInsnNode.name}' should go after LDC but $prev" }
                     val prevPev = methodInsnNode.previous.previous

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

@@ -5,80 +5,13 @@
 
 package org.jetbrains.kotlin.codegen.inline
 
+import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethods
 import org.jetbrains.kotlin.codegen.optimization.common.InsnSequence
 import org.jetbrains.kotlin.codegen.optimization.common.isMeaningful
-import org.jetbrains.kotlin.codegen.optimization.fixStack.top
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterSignature
-import org.jetbrains.kotlin.utils.SmartSet
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.tree.*
-import org.jetbrains.org.objectweb.asm.tree.analysis.Frame
+import org.jetbrains.org.objectweb.asm.tree.analysis.*
-import org.jetbrains.org.objectweb.asm.tree.analysis.SourceValue
-
-fun MethodInliner.getFunctionalArgumentIfExistsAndMarkInstructions(
-    sourceValue: SourceValue,
-    processSwap: Boolean,
-    insnList: InsnList,
-    frames: Array<Frame<SourceValue>?>,
-    toDelete: MutableSet<AbstractInsnNode>
-): FunctionalArgument? {
-    val toDeleteInner = SmartSet.create<AbstractInsnNode>()
-
-    val functionalArgumentSet = SmartSet.create<FunctionalArgument?>()
-    sourceValue.insns.mapTo(functionalArgumentSet) {
-        getFunctionalArgumentIfExistsAndMarkInstructions(it, processSwap, insnList, frames, toDeleteInner)
-    }
-
-    return functionalArgumentSet.singleOrNull()?.also {
-        if (it is LambdaInfo) {
-            toDelete.addAll(toDeleteInner)
-        }
-    }
-}
-
-private fun SourceValue.singleOrNullInsn() = insns.singleOrNull()
-
-private fun MethodInliner.getFunctionalArgumentIfExistsAndMarkInstructions(
-    insnNode: AbstractInsnNode?,
-    processSwap: Boolean,
-    insnList: InsnList,
-    frames: Array<Frame<SourceValue>?>,
-    toDelete: MutableSet<AbstractInsnNode>
-): FunctionalArgument? {
-    if (insnNode == null) return null
-
-    getFunctionalArgumentIfExists(insnNode)?.let {
-        //delete lambda aload instruction
-        toDelete.add(insnNode)
-        return it
-    }
-
-    if (insnNode is VarInsnNode && insnNode.opcode == Opcodes.ALOAD) {
-        val varIndex = insnNode.`var`
-        val localFrame = frames[insnList.indexOf(insnNode)] ?: return null
-        val storeIns = localFrame.getLocal(varIndex).singleOrNullInsn()
-        if (storeIns is VarInsnNode && storeIns.getOpcode() == Opcodes.ASTORE) {
-            val frame = frames[insnList.indexOf(storeIns)] ?: return null
-            val topOfStack = frame.top()!!
-            getFunctionalArgumentIfExistsAndMarkInstructions(topOfStack, processSwap, insnList, frames, toDelete)?.let {
-                //remove intermediate lambda astore, aload instruction: see 'complexStack/simple.1.kt' test
-                toDelete.add(storeIns)
-                toDelete.add(insnNode)
-                return it
-            }
-        }
-    } else if (processSwap && insnNode.opcode == Opcodes.SWAP) {
-        val swapFrame = frames[insnList.indexOf(insnNode)] ?: return null
-        val dispatchReceiver = swapFrame.top()!!
-        getFunctionalArgumentIfExistsAndMarkInstructions(dispatchReceiver, false, insnList, frames, toDelete)?.let {
-            //remove swap instruction (dispatch receiver would be deleted on recursion call): see 'complexStack/simpleExtension.1.kt' test
-            toDelete.add(insnNode)
-            return it
-        }
-    }
-
-    return null
-}
 
 fun parameterOffsets(isStatic: Boolean, valueParameters: List<JvmMethodParameterSignature>): Array<Int> {
     var nextOffset = if (isStatic) 0 else 1
@@ -121,4 +54,90 @@ fun AbstractInsnNode.getNextMeaningful(): AbstractInsnNode? {
         result = result.next
     }
     return result
+}
+
+// Interpreter, that analyzes only functional arguments, to replace SourceInterpreter, since SourceInterpreter' merge has O(N²) complexity
+
+internal open class FunctionalArgumentValue(
+    val functionalArgument: FunctionalArgument, basicValue: BasicValue?
+) : BasicValue(basicValue?.type) {
+    override fun toString(): String = "$functionalArgument"
+}
+
+val BasicValue?.functionalArgument
+    get() = (this as? FunctionalArgumentValue)?.functionalArgument
+
+private class FunctionalArgumentInterpreter(
+    private val inliner: MethodInliner, private val toDelete: MutableSet<AbstractInsnNode>
+) : BasicInterpreter(Opcodes.ASM5) {
+
+    override fun unaryOperation(insn: AbstractInsnNode, value: BasicValue): BasicValue? =
+        markInstructionIfNeeded(insn, super.unaryOperation(insn, value))
+
+    override fun copyOperation(insn: AbstractInsnNode, value: BasicValue?): BasicValue? {
+        val basicValue = super.copyOperation(insn, value)
+        // Paramter checks are processed separately
+        if (insn.next?.opcode == Opcodes.LDC && isParameterNullabilityCheck(insn.next?.next)) {
+            return basicValue
+        }
+        if (value.functionalArgument is LambdaInfo) {
+            // SWAP and ASTORE
+            toDelete.add(insn)
+            return value
+        }
+        return markInstructionIfNeeded(insn, basicValue)
+    }
+
+    private fun markInstructionIfNeeded(
+        insn: AbstractInsnNode,
+        basicValue: BasicValue?
+    ): BasicValue? {
+        val functionalArgument = inliner.getFunctionalArgumentIfExists(insn)
+        return if (functionalArgument != null) {
+            if (functionalArgument is LambdaInfo) {
+                toDelete.add(insn)
+            }
+            FunctionalArgumentValue(functionalArgument, basicValue)
+        } else basicValue
+    }
+
+    override fun newOperation(insn: AbstractInsnNode): BasicValue? =
+        markInstructionIfNeeded(insn, super.newOperation(insn))
+
+    override fun merge(v: BasicValue?, w: BasicValue?): BasicValue? =
+        if (v is FunctionalArgumentValue && w is FunctionalArgumentValue && v.functionalArgument == w.functionalArgument) v
+        else super.merge(v, w)
+}
+
+fun isParameterNullabilityCheck(methodInsnNode: AbstractInsnNode?): Boolean =
+    methodInsnNode is MethodInsnNode && methodInsnNode.owner == IntrinsicMethods.INTRINSICS_CLASS_NAME &&
+            (methodInsnNode.name == "checkParameterIsNotNull" || methodInsnNode.name == "checkNotNullParameter")
+
+internal fun analyzeMethodWithFunctionalArgumentInterpreter(
+    inliner: MethodInliner,
+    node: MethodNode,
+    toDelete: MutableSet<AbstractInsnNode>
+): Array<out Frame<BasicValue>?> {
+    val analyzer = object : Analyzer<BasicValue>(FunctionalArgumentInterpreter(inliner, toDelete)) {
+        override fun newFrame(nLocals: Int, nStack: Int): Frame<BasicValue> {
+
+            return object : Frame<BasicValue>(nLocals, nStack) {
+                @Throws(AnalyzerException::class)
+                override fun execute(insn: AbstractInsnNode, interpreter: Interpreter<BasicValue>) {
+                    if (node.name == "waitForEx\$default" && node.instructions.indexOf(insn) == 18) {
+                        {}()
+                    }
+                    // This can be a void non-local return from a non-void method; Frame#execute would throw and do nothing else.
+                    if (insn.opcode == Opcodes.RETURN) return
+                    super.execute(insn, interpreter)
+                }
+            }
+        }
+    }
+
+    try {
+        return analyzer.analyze("fake", node)
+    } catch (e: AnalyzerException) {
+        throw RuntimeException(e)
+    }
 }

compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn.kt

@@ -0,0 +1,15 @@
+// FILE: 1.kt
+fun check1() = true
+
+fun check2() = false
+
+inline fun inlineMe1(fn: (String, String) -> String): String {
+    return fn(if (check1()) return "O" else "1", return "2")
+}
+
+inline fun inlineMe2(fn: (String) -> String): String {
+    return fn(if (check2()) return "3" else "K")
+}
+
+// FILE: 2.kt
+fun box() = inlineMe1 { _, _ -> "FAIL1" } + inlineMe2 { it }

compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn2.kt

@@ -0,0 +1,10 @@
+// FILE: 1.kt
+inline fun alwaysOk(s: String, fn: (String) -> String): String {
+    return fn(try {return "OK"} catch (e: Exception) {
+        "fail1"
+    }) + fn(try {return "FF"} catch (e: Exception) {
+        "fail2"
+    })
+}
+// FILE: 2.kt
+fun box() = alwaysOk("what?") { it }

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

@@ -3696,6 +3696,16 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/nonLocalReturn3.kt");
         }
 
+        @TestMetadata("poppedLocalReturn.kt")
+        public void testPoppedLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn.kt");
+        }
+
+        @TestMetadata("poppedLocalReturn2.kt")
+        public void testPoppedLocalReturn2() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn2.kt");
+        }
+
         @TestMetadata("returnLong.kt")
         public void testReturnLong() throws Exception {
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/returnLong.kt");

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

@@ -3696,6 +3696,16 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/nonLocalReturn3.kt");
         }
 
+        @TestMetadata("poppedLocalReturn.kt")
+        public void testPoppedLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn.kt");
+        }
+
+        @TestMetadata("poppedLocalReturn2.kt")
+        public void testPoppedLocalReturn2() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn2.kt");
+        }
+
         @TestMetadata("returnLong.kt")
         public void testReturnLong() throws Exception {
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/returnLong.kt");

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

@@ -3696,6 +3696,16 @@ public class IrBlackBoxInlineCodegenTestGenerated extends AbstractIrBlackBoxInli
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/nonLocalReturn3.kt");
         }
 
+        @TestMetadata("poppedLocalReturn.kt")
+        public void testPoppedLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn.kt");
+        }
+
+        @TestMetadata("poppedLocalReturn2.kt")
+        public void testPoppedLocalReturn2() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn2.kt");
+        }
+
         @TestMetadata("returnLong.kt")
         public void testReturnLong() throws Exception {
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/returnLong.kt");

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

@@ -3696,6 +3696,16 @@ public class IrCompileKotlinAgainstInlineKotlinTestGenerated extends AbstractIrC
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/nonLocalReturn3.kt");
         }
 
+        @TestMetadata("poppedLocalReturn.kt")
+        public void testPoppedLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn.kt");
+        }
+
+        @TestMetadata("poppedLocalReturn2.kt")
+        public void testPoppedLocalReturn2() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn2.kt");
+        }
+
         @TestMetadata("returnLong.kt")
         public void testReturnLong() throws Exception {
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/returnLong.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenInlineTestGenerated.java

@@ -3296,6 +3296,16 @@ public class IrJsCodegenInlineTestGenerated extends AbstractIrJsCodegenInlineTes
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/nonLocalReturn3.kt");
         }
 
+        @TestMetadata("poppedLocalReturn.kt")
+        public void testPoppedLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn.kt");
+        }
+
+        @TestMetadata("poppedLocalReturn2.kt")
+        public void testPoppedLocalReturn2() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn2.kt");
+        }
+
         @TestMetadata("returnLong.kt")
         public void testReturnLong() throws Exception {
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/returnLong.kt");

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

@@ -3296,6 +3296,16 @@ public class JsCodegenInlineTestGenerated extends AbstractJsCodegenInlineTest {
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/nonLocalReturn3.kt");
         }
 
+        @TestMetadata("poppedLocalReturn.kt")
+        public void testPoppedLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn.kt");
+        }
+
+        @TestMetadata("poppedLocalReturn2.kt")
+        public void testPoppedLocalReturn2() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/stackOnReturn/poppedLocalReturn2.kt");
+        }
+
         @TestMetadata("returnLong.kt")
         public void testReturnLong() throws Exception {
             runTest("compiler/testData/codegen/boxInline/stackOnReturn/returnLong.kt");