Fix KT-8608: Compiler crashes with assertion Restore stack is unavailable

- fix SAVE_STACK_BEFORE_TRY insertion:
  TRYCATCHBLOCK LA, LB, LC
  LA
    NOP
    try_body
  LB
    ...
  LC
    handler_body
should be transformed into:
  LA
    {SAVE_STACK_BEFORE_TRY}
  LA' // new TCB start label
    NOP
    try_body
  LB
    ...
  LC
    handler_body
with all TCBs start labels remapped

- properly wrap exceptions from MandatoryMethodTransformer

 #KT-8608 Fixed

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

@@ -374,7 +374,12 @@ public class MethodInliner {
     protected MethodNode markPlacesForInlineAndRemoveInlinable(@NotNull MethodNode node, int finallyDeepShift) {
         node = prepareNode(node, finallyDeepShift);
 
-        MandatoryMethodTransformer.INSTANCE$.transform("fake", node);
+        try {
+            MandatoryMethodTransformer.INSTANCE$.transform("fake", node);
+        }
+        catch (Throwable e) {
+            throw wrapException(e, node, "couldn't inline method call");
+        }
 
         Analyzer<SourceValue> analyzer = new Analyzer<SourceValue>(new SourceInterpreter()) {
             @NotNull
@@ -689,7 +694,7 @@ public class MethodInliner {
     }
 
     @NotNull
-    public RuntimeException wrapException(@NotNull Exception originalException, @NotNull MethodNode node, @NotNull String errorSuffix) {
+    public RuntimeException wrapException(@NotNull Throwable originalException, @NotNull MethodNode node, @NotNull String errorSuffix) {
         if (originalException instanceof InlineException) {
             return new InlineException(errorPrefix + ": " + errorSuffix, originalException);
         }

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/fixStack/AnalyzeTryCatchBlocks.kt

@@ -21,8 +21,10 @@ import org.jetbrains.kotlin.codegen.optimization.common.findNextOrNull
 import org.jetbrains.kotlin.codegen.optimization.common.hasOpcode
 import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
 import org.jetbrains.kotlin.codegen.pseudoInsns.PseudoInsn
+import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.tree.*
 import org.jetbrains.org.objectweb.asm.util.Printer
+import java.util.*
 
 private class DecompiledTryDescriptor(val tryStartLabel: LabelNode) {
     var defaultHandlerTcb : TryCatchBlockNode? = null
@@ -41,6 +43,71 @@ internal fun insertTryCatchBlocksMarkers(methodNode: MethodNode) {
     val decompiledTryDescriptorForStart = linkedMapOf<LabelNode, DecompiledTryDescriptor>()
     val decompiledTryDescriptorForHandler = hashMapOf<LabelNode, DecompiledTryDescriptor>()
 
+    collectDecompiledTryDescriptors(decompiledTryDescriptorForStart, decompiledTryDescriptorForHandler, methodNode)
+
+
+    val newTryStartLabels = hashMapOf<LabelNode, LabelNode>()
+
+    insertSaveRestoreStackMarkers(decompiledTryDescriptorForStart, methodNode, newTryStartLabels)
+
+    transformTryCatchBlocks(methodNode, newTryStartLabels)
+}
+
+private fun transformTryCatchBlocks(methodNode: MethodNode, newTryStartLabels: HashMap<LabelNode, LabelNode>) {
+    methodNode.tryCatchBlocks = methodNode.tryCatchBlocks.map { tcb ->
+        val newTryStartLabel = newTryStartLabels[tcb.start]
+        if (newTryStartLabel == null)
+            tcb
+        else
+            TryCatchBlockNode(newTryStartLabel, tcb.end, tcb.handler, tcb.type)
+    }
+}
+
+private fun insertSaveRestoreStackMarkers(
+        decompiledTryDescriptorForStart: Map<LabelNode, DecompiledTryDescriptor>,
+        methodNode: MethodNode,
+        newTryStartLabels: MutableMap<LabelNode, LabelNode>
+) {
+    val doneTryStartLabels = hashSetOf<LabelNode>()
+    val doneHandlerLabels = hashSetOf<LabelNode>()
+
+    for (decompiledTryDescriptor in decompiledTryDescriptorForStart.values()) {
+        with(decompiledTryDescriptor) {
+            if (!doneTryStartLabels.contains(tryStartLabel)) {
+                doneTryStartLabels.add(tryStartLabel)
+
+                val nopNode = tryStartLabel.findNextOrNull { it.hasOpcode() }!!
+                assert(nopNode.getOpcode() == Opcodes.NOP,
+                       "${methodNode.instructions.indexOf(nopNode)}: try block should start with NOP")
+
+                val newTryStartLabel = LabelNode(Label())
+                newTryStartLabels[tryStartLabel] = newTryStartLabel
+
+                methodNode.instructions.insertBefore(nopNode, PseudoInsn.SAVE_STACK_BEFORE_TRY.createInsnNode())
+                methodNode.instructions.insertBefore(nopNode, newTryStartLabel)
+                methodNode.instructions.insert(nopNode, PseudoInsn.RESTORE_STACK_IN_TRY_CATCH.createInsnNode())
+            }
+
+            for (handlerStartLabel in handlerStartLabels) {
+                if (!doneHandlerLabels.contains(handlerStartLabel)) {
+                    doneHandlerLabels.add(handlerStartLabel)
+
+                    val storeNode = handlerStartLabel.findNextOrNull { it.hasOpcode() }!!
+                    assert(storeNode.getOpcode() == Opcodes.ASTORE,
+                           "${methodNode.instructions.indexOf(storeNode)}: handler should start with ASTORE")
+
+                    methodNode.instructions.insert(storeNode, PseudoInsn.RESTORE_STACK_IN_TRY_CATCH.createInsnNode())
+                }
+            }
+        }
+    }
+}
+
+private fun collectDecompiledTryDescriptors(
+        decompiledTryDescriptorForStart: MutableMap<LabelNode, DecompiledTryDescriptor>,
+        decompiledTryDescriptorForHandler: MutableMap<LabelNode, DecompiledTryDescriptor>,
+        methodNode: MethodNode
+) {
     for (tcb in methodNode.tryCatchBlocks) {
         if (tcb.isDefaultHandlerNode()) {
             assert(decompiledTryDescriptorForHandler.containsKey(tcb.start),
@@ -64,34 +131,4 @@ internal fun insertTryCatchBlocksMarkers(methodNode: MethodNode) {
             }
         }
     }
-
-    val doneTryStartLabels = hashSetOf<LabelNode>()
-    val doneHandlerLabels = hashSetOf<LabelNode>()
-
-    for (decompiledTryDescriptor in decompiledTryDescriptorForStart.values()) {
-        with(decompiledTryDescriptor) {
-            if (!doneTryStartLabels.contains(tryStartLabel)) {
-                doneTryStartLabels.add(tryStartLabel)
-
-                val nopNode = tryStartLabel.findNextOrNull { it.hasOpcode() }!!
-                assert(nopNode.getOpcode() == Opcodes.NOP,
-                       "${methodNode.instructions.indexOf(nopNode)}: try block should start with NOP")
-
-                methodNode.instructions.insertBefore(tryStartLabel, PseudoInsn.SAVE_STACK_BEFORE_TRY.createInsnNode())
-                methodNode.instructions.insert(nopNode, PseudoInsn.RESTORE_STACK_IN_TRY_CATCH.createInsnNode())
-            }
-
-            for (handlerStartLabel in handlerStartLabels) {
-                if (!doneHandlerLabels.contains(handlerStartLabel)) {
-                    doneHandlerLabels.add(handlerStartLabel)
-
-                    val storeNode = handlerStartLabel.findNextOrNull { it.hasOpcode() }!!
-                    assert(storeNode.getOpcode() == Opcodes.ASTORE,
-                           "${methodNode.instructions.indexOf(storeNode)}: handler should start with ASTORE")
-
-                    methodNode.instructions.insert(storeNode, PseudoInsn.RESTORE_STACK_IN_TRY_CATCH.createInsnNode())
-                }
-            }
-        }
-    }
 }

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/fixStack/FixStackMethodTransformer.kt

@@ -148,14 +148,8 @@ public object FixStackMethodTransformer : MethodTransformer() {
             marker: AbstractInsnNode,
             localVariablesManager: LocalVariablesManager
     ) {
-        val savedStackDescriptor = localVariablesManager.getSavedStackDescriptorOrNull(marker)
+        val savedStackDescriptor = localVariablesManager.getSavedStackDescriptor(marker)
-        if (savedStackDescriptor != null) {
+        actions.add({ restoreStack(methodNode, marker, savedStackDescriptor) })
-            actions.add({ restoreStack(methodNode, marker, savedStackDescriptor) })
-        }
-        else {
-            // marker is dead code
-            actions.add({ methodNode.instructions.remove(marker) })
-        }
         localVariablesManager.markRestoreStackMarkerEmitted(marker)
     }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/fixStack/LocalVariablesManager.kt

@@ -52,7 +52,7 @@ internal class LocalVariablesManager(val context: FixStackContext, val methodNod
         return savedStackDescriptor
     }
 
-    fun getSavedStackDescriptorOrNull(restoreStackMarker: AbstractInsnNode): SavedStackDescriptor {
+    fun getSavedStackDescriptor(restoreStackMarker: AbstractInsnNode): SavedStackDescriptor {
         val saveStackMarker = context.saveStackMarkerForRestoreMarker[restoreStackMarker]
         return allocatedHandles[saveStackMarker]!!.savedStackDescriptor
     }

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

@@ -0,0 +1,26 @@
+inline fun catchAll(x: String, block: () -> Unit): String {
+    try {
+        block()
+    } catch (e: Throwable) {
+    }
+    return x
+}
+
+inline fun tryTwice(block: () -> Unit) {
+    try {
+        block()
+        try {
+            block()
+        } catch (e: Exception) {
+        }
+    } catch (e: Exception) {
+    }
+}
+
+fun box(): String {
+    return catchAll("OK") {
+        tryTwice {
+            throw Exception()
+        }
+    }
+}

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

@@ -0,0 +1,30 @@
+interface Callable {
+    fun call(b: Boolean)
+}
+
+inline fun run(f: () -> Unit) { f() }
+
+class A {
+    fun foo(): String {
+        run {
+            val x = object : Callable {
+                override fun call(b: Boolean) {
+                    if (b) {
+                        x()
+                    } else {
+                        try {
+                            x()
+                        } catch(t: Throwable) {
+                        }
+                    }
+                }
+            }
+        }
+        return "OK"
+    }
+
+    private fun x() {}
+}
+
+fun box(): String =
+        A().foo()

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

@@ -2322,6 +2322,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("deadTryCatch.kt")
+            public void testDeadTryCatch() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/deadTryCatch.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("differentTypes.kt")
             public void testDifferentTypes() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/differentTypes.kt");
@@ -2358,6 +2364,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("kt8608.kt")
+            public void testKt8608() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/kt8608.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("multipleCatchBlocks.kt")
             public void testMultipleCatchBlocks() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/tryCatchInExpressions/multipleCatchBlocks.kt");