Traverse multiple store-load chains for inlined lambda parameters

When a try-catch expression is passed as an argument to the inline
lambda parameter, lambda variable on stack is spilled and restored in
several different locations (1 for try-block, 1 for each catch-blocks).
So it's possible that lambda to be invoked comes from multiple loads,
all of which should have the same "root" lambda parameter.

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

@@ -456,7 +456,7 @@ public class MethodInliner {
                         SourceValue sourceValue = frame.getStack(firstParameterIndex);
 
                         LambdaInfo lambdaInfo = MethodInlinerUtilKt.getLambdaIfExistsAndMarkInstructions(
-                                this, MethodInlinerUtilKt.singleOrNullInsn(sourceValue), true, instructions, sources, toDelete
+                                this, sourceValue, true, instructions, sources, toDelete
                         );
 
                         invokeCalls.add(new InvokeCall(lambdaInfo, currentFinallyDeep));
@@ -469,7 +469,7 @@ public class MethodInliner {
                         for (int i = 0; i < paramCount; i++) {
                             SourceValue sourceValue = frame.getStack(firstParameterIndex + i);
                             LambdaInfo lambdaInfo = MethodInlinerUtilKt.getLambdaIfExistsAndMarkInstructions(
-                                    this, MethodInlinerUtilKt.singleOrNullInsn(sourceValue), false, instructions, sources, toDelete
+                                    this, sourceValue, false, instructions, sources, toDelete
                             );
                             if (lambdaInfo != null) {
                                 lambdaMapping.put(offset, lambdaInfo);

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

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.codegen.inline
 
 import org.jetbrains.kotlin.codegen.optimization.fixStack.top
+import org.jetbrains.kotlin.utils.SmartSet
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
 import org.jetbrains.org.objectweb.asm.tree.InsnList
@@ -25,6 +26,25 @@ import org.jetbrains.org.objectweb.asm.tree.analysis.Frame
 import org.jetbrains.org.objectweb.asm.tree.analysis.SourceValue
 
 fun MethodInliner.getLambdaIfExistsAndMarkInstructions(
+        sourceValue: SourceValue,
+        processSwap: Boolean,
+        insnList: InsnList,
+        frames: Array<Frame<SourceValue>?>,
+        toDelete: MutableSet<AbstractInsnNode>
+): LambdaInfo? {
+    val toDeleteInner = SmartSet.create<AbstractInsnNode>()
+
+    val lambdaSet = SmartSet.create<LambdaInfo?>()
+    sourceValue.insns.mapTo(lambdaSet) {
+        getLambdaIfExistsAndMarkInstructions(it, processSwap, insnList, frames, toDeleteInner)
+    }
+
+    return lambdaSet.singleOrNull()?.also {
+        toDelete.addAll(toDeleteInner)
+    }
+}
+
+private fun MethodInliner.getLambdaIfExistsAndMarkInstructions(
         insnNode: AbstractInsnNode?,
         processSwap: Boolean,
         insnList: InsnList,
@@ -45,7 +65,7 @@ fun MethodInliner.getLambdaIfExistsAndMarkInstructions(
         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()!!.singleOrNullInsn()
+            val topOfStack = frame.top()!!
             getLambdaIfExistsAndMarkInstructions(topOfStack, processSwap, insnList, frames, toDelete)?.let {
                 //remove intermediate lambda astore, aload instruction: see 'complexStack/simple.1.kt' test
                 toDelete.add(storeIns)
@@ -56,7 +76,7 @@ fun MethodInliner.getLambdaIfExistsAndMarkInstructions(
     }
     else if (processSwap && insnNode.opcode == Opcodes.SWAP) {
         val swapFrame = frames[insnList.indexOf(insnNode)] ?: return null
-        val dispatchReceiver = swapFrame.top()!!.singleOrNullInsn()
+        val dispatchReceiver = swapFrame.top()!!
         getLambdaIfExistsAndMarkInstructions(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)

compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572.kt

@@ -0,0 +1,6 @@
+fun zap(s: String) = s
+
+inline fun tryZap(string: String, fn: (String) -> String) =
+        fn(try { zap(string) } catch (e: Exception) { "" })
+
+fun box(): String = tryZap("OK") { it }

compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2.kt

@@ -0,0 +1,9 @@
+fun zap(s: String) = s
+
+inline fun tryZap(s1: String, s2: String, fn: (String, String) -> String) =
+        fn(
+                try { zap(s1) } catch (e: Exception) { "" },
+                try { zap(s2) } catch (e: Exception) { "" }
+        )
+
+fun box(): String = tryZap("O", "K") { a, b -> a + b }

compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2_ext.kt

@@ -0,0 +1,9 @@
+fun zap(s: String) = s
+
+inline fun tryZap(s1: String, s2: String, fn: String.(String) -> String) =
+        fn(
+                try { zap(s1) } catch (e: Exception) { "" },
+                try { zap(s2) } catch (e: Exception) { "" }
+        )
+
+fun box(): String = tryZap("O", "K") { this + it }

compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_ext.kt

@@ -0,0 +1,6 @@
+fun zap(s: String) = s
+
+inline fun tryZap(string: String, fn: String.() -> String) =
+        fn(try { zap(string) } catch (e: Exception) { "" })
+
+fun box(): String = tryZap("OK") { this }

compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_nested.kt

@@ -0,0 +1,13 @@
+fun zap(s: String) = s
+
+inline fun tryZap(string: String, fn: (String) -> String) =
+        fn(
+                try {
+                    try {
+                        zap(string)
+                    }
+                    catch (e: Exception) { "" }
+                } catch (e: Exception) { "" }
+        )
+
+fun box(): String = tryZap("OK") { it }

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

@@ -4780,6 +4780,36 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 doTest(fileName);
             }
 
+            @TestMetadata("kt17572.kt")
+            public void testKt17572() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_2.kt")
+            public void testKt17572_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_2_ext.kt")
+            public void testKt17572_2_ext() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2_ext.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_ext.kt")
+            public void testKt17572_ext() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_ext.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_nested.kt")
+            public void testKt17572_nested() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_nested.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("kt8608.kt")
             public void testKt8608() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt8608.kt");

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

@@ -4780,6 +4780,36 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("kt17572.kt")
+            public void testKt17572() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_2.kt")
+            public void testKt17572_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_2_ext.kt")
+            public void testKt17572_2_ext() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2_ext.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_ext.kt")
+            public void testKt17572_ext() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_ext.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_nested.kt")
+            public void testKt17572_nested() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_nested.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("kt8608.kt")
             public void testKt8608() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt8608.kt");

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

@@ -4780,6 +4780,36 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 doTest(fileName);
             }
 
+            @TestMetadata("kt17572.kt")
+            public void testKt17572() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_2.kt")
+            public void testKt17572_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_2_ext.kt")
+            public void testKt17572_2_ext() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2_ext.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_ext.kt")
+            public void testKt17572_ext() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_ext.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_nested.kt")
+            public void testKt17572_nested() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_nested.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("kt8608.kt")
             public void testKt8608() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt8608.kt");

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

@@ -5495,6 +5495,36 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("kt17572.kt")
+            public void testKt17572() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_2.kt")
+            public void testKt17572_2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_2_ext.kt")
+            public void testKt17572_2_ext() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_2_ext.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_ext.kt")
+            public void testKt17572_ext() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_ext.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt17572_nested.kt")
+            public void testKt17572_nested() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt17572_nested.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("kt8608.kt")
             public void testKt8608() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt8608.kt");

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

@@ -59,4 +59,5 @@ public class OutputPrefixPostfixTestGenerated extends AbstractOutputPrefixPostfi
         String fileName = KotlinTestUtils.navigationMetadata("js/js.translator/testData/outputPrefixPostfix/simpleWithPrefixAndPostfix.kt");
         doTest(fileName);
     }
+
 }