Minor. Simplify collecting save/restore markers for FixStackMethodTransformer

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

@@ -19,14 +19,16 @@ package org.jetbrains.kotlin.codegen.optimization.fixStack
 import com.sun.xml.internal.ws.org.objectweb.asm.Opcodes
 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.tree.AbstractInsnNode
-import org.jetbrains.org.objectweb.asm.util.Printer
+import org.jetbrains.org.objectweb.asm.tree.LabelNode
+import org.jetbrains.org.objectweb.asm.tree.MethodNode
+import org.jetbrains.org.objectweb.asm.tree.TryCatchBlockNode
 import java.util.*
 
 private class DecompiledTryDescriptor(val tryStartLabel: LabelNode) {
+    // Only used for assertions
     var defaultHandlerTcb : TryCatchBlockNode? = null
     val handlerStartLabels = hashSetOf<LabelNode>()
 }
@@ -37,20 +39,17 @@ private fun TryCatchBlockNode.isDefaultHandlerNode(): Boolean =
 private fun MethodNode.debugString(tcb: TryCatchBlockNode): String =
         "TCB<${instructions.indexOf(tcb.start)}, ${instructions.indexOf(tcb.end)}, ${instructions.indexOf(tcb.handler)}>"
 
-internal fun insertTryCatchBlocksMarkers(methodNode: MethodNode) {
+internal fun insertTryCatchBlocksMarkers(methodNode: MethodNode): Map<AbstractInsnNode, AbstractInsnNode> {
-    if (methodNode.tryCatchBlocks.isEmpty()) return
+    if (methodNode.tryCatchBlocks.isEmpty()) return emptyMap()
-
-    val decompiledTryDescriptorForStart = linkedMapOf<LabelNode, DecompiledTryDescriptor>()
-    val decompiledTryDescriptorForHandler = hashMapOf<LabelNode, DecompiledTryDescriptor>()
-
-    collectDecompiledTryDescriptors(decompiledTryDescriptorForStart, decompiledTryDescriptorForHandler, methodNode)
 
+    val decompiledTryDescriptorForStart = collectDecompiledTryDescriptors(methodNode)
 
     val newTryStartLabels = hashMapOf<LabelNode, LabelNode>()
-
+    val restoreStackToSaveStackMarker = insertSaveRestoreStackMarkers(decompiledTryDescriptorForStart, methodNode, newTryStartLabels)
-    insertSaveRestoreStackMarkers(decompiledTryDescriptorForStart, methodNode, newTryStartLabels)
 
     transformTryCatchBlocks(methodNode, newTryStartLabels)
+
+    return restoreStackToSaveStackMarker
 }
 
 private fun transformTryCatchBlocks(methodNode: MethodNode, newTryStartLabels: HashMap<LabelNode, LabelNode>) {
@@ -67,24 +66,34 @@ private fun insertSaveRestoreStackMarkers(
         decompiledTryDescriptorForStart: Map<LabelNode, DecompiledTryDescriptor>,
         methodNode: MethodNode,
         newTryStartLabels: MutableMap<LabelNode, LabelNode>
-) {
+): Map<AbstractInsnNode, AbstractInsnNode> {
-    val doneTryStartLabels = hashSetOf<LabelNode>()
+    val restoreStackToSaveMarker = hashMapOf<AbstractInsnNode, AbstractInsnNode>()
+    val saveStackMarkerByTryLabel = hashMapOf<LabelNode, AbstractInsnNode>()
     val doneHandlerLabels = hashSetOf<LabelNode>()
 
     for (decompiledTryDescriptor in decompiledTryDescriptorForStart.values) {
         with(decompiledTryDescriptor) {
-            if (!doneTryStartLabels.contains(tryStartLabel)) {
+            val saveStackMarker: AbstractInsnNode
-                doneTryStartLabels.add(tryStartLabel)
 
+            if (tryStartLabel !in saveStackMarkerByTryLabel) {
                 val nopNode = tryStartLabel.findNextOrNull { it.hasOpcode() }!!
-                assert(nopNode.getOpcode() == Opcodes.NOP) { "${methodNode.instructions.indexOf(nopNode)}: try block should start with NOP" }
+                assert(nopNode.opcode == 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())
+                saveStackMarker = PseudoInsn.SAVE_STACK_BEFORE_TRY.createInsnNode()
+                val restoreStackMarker = PseudoInsn.RESTORE_STACK_IN_TRY_CATCH.createInsnNode()
+
+                saveStackMarkerByTryLabel[tryStartLabel] = saveStackMarker
+                restoreStackToSaveMarker[restoreStackMarker] = saveStackMarker
+
+                methodNode.instructions.insertBefore(nopNode, saveStackMarker)
                 methodNode.instructions.insertBefore(nopNode, newTryStartLabel)
-                methodNode.instructions.insert(nopNode, PseudoInsn.RESTORE_STACK_IN_TRY_CATCH.createInsnNode())
+                methodNode.instructions.insert(nopNode, restoreStackMarker)
+            }
+            else {
+                saveStackMarker = saveStackMarkerByTryLabel[tryStartLabel]!!
             }
 
             for (handlerStartLabel in handlerStartLabels) {
@@ -92,20 +101,22 @@ private fun insertSaveRestoreStackMarkers(
                     doneHandlerLabels.add(handlerStartLabel)
 
                     val storeNode = handlerStartLabel.findNextOrNull { it.hasOpcode() }!!
-                    assert(storeNode.getOpcode() == Opcodes.ASTORE) { "${methodNode.instructions.indexOf(storeNode)}: handler should start with ASTORE" }
+                    assert(storeNode.opcode == Opcodes.ASTORE) { "${methodNode.instructions.indexOf(storeNode)}: handler should start with ASTORE" }
 
-                    methodNode.instructions.insert(storeNode, PseudoInsn.RESTORE_STACK_IN_TRY_CATCH.createInsnNode())
+                    val restoreStackMarker = PseudoInsn.RESTORE_STACK_IN_TRY_CATCH.createInsnNode()
+                    restoreStackToSaveMarker[restoreStackMarker] = saveStackMarker
+                    methodNode.instructions.insert(storeNode, restoreStackMarker)
                 }
             }
         }
     }
+
+    return restoreStackToSaveMarker
 }
 
-private fun collectDecompiledTryDescriptors(
+private fun collectDecompiledTryDescriptors(methodNode: MethodNode): Map<LabelNode, DecompiledTryDescriptor> {
-        decompiledTryDescriptorForStart: MutableMap<LabelNode, DecompiledTryDescriptor>,
+    val decompiledTryDescriptorForStart: MutableMap<LabelNode, DecompiledTryDescriptor> = hashMapOf()
-        decompiledTryDescriptorForHandler: MutableMap<LabelNode, DecompiledTryDescriptor>,
+    val decompiledTryDescriptorForHandler: MutableMap<LabelNode, DecompiledTryDescriptor> = hashMapOf()
-        methodNode: MethodNode
-) {
     for (tcb in methodNode.tryCatchBlocks) {
         if (tcb.isDefaultHandlerNode()) {
             assert(decompiledTryDescriptorForHandler.containsKey(tcb.start)) { "${methodNode.debugString(tcb)}: default handler should occur after some regular handler" }
@@ -128,4 +139,6 @@ private fun collectDecompiledTryDescriptors(
             }
         }
     }
+
+    return decompiledTryDescriptorForStart
 }

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

@@ -17,18 +17,14 @@
 package org.jetbrains.kotlin.codegen.optimization.fixStack
 
 import com.intellij.util.SmartList
-import com.intellij.util.containers.SmartHashSet
 import com.intellij.util.containers.Stack
 import org.jetbrains.kotlin.codegen.inline.InlineCodegenUtil
 import org.jetbrains.kotlin.codegen.optimization.common.InsnSequence
-import org.jetbrains.kotlin.codegen.optimization.common.findPreviousOrNull
-import org.jetbrains.kotlin.codegen.optimization.common.hasOpcode
 import org.jetbrains.kotlin.codegen.pseudoInsns.PseudoInsn
 import org.jetbrains.kotlin.codegen.pseudoInsns.parsePseudoInsnOrNull
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
 import org.jetbrains.org.objectweb.asm.tree.JumpInsnNode
-import org.jetbrains.org.objectweb.asm.tree.LabelNode
 import org.jetbrains.org.objectweb.asm.tree.MethodNode
 
 internal class FixStackContext(val methodNode: MethodNode) {
@@ -36,15 +32,19 @@ internal class FixStackContext(val methodNode: MethodNode) {
     val fakeAlwaysTrueIfeqMarkers = arrayListOf<AbstractInsnNode>()
     val fakeAlwaysFalseIfeqMarkers = arrayListOf<AbstractInsnNode>()
 
-    val saveStackNodesForTryStartLabel = hashMapOf<LabelNode, AbstractInsnNode>()
+    val isThereAnyTryCatch: Boolean
-    val saveStackMarkerForRestoreMarker = hashMapOf<AbstractInsnNode, AbstractInsnNode>()
+    val saveStackMarkerForRestoreMarker: Map<AbstractInsnNode, AbstractInsnNode>
     val restoreStackMarkersForSaveMarker = hashMapOf<AbstractInsnNode, MutableList<AbstractInsnNode>>()
 
     val openingInlineMethodMarker = hashMapOf<AbstractInsnNode, AbstractInsnNode>()
     var consistentInlineMarkers: Boolean = true; private set
 
     init {
-        insertTryCatchBlocksMarkers(methodNode)
+        saveStackMarkerForRestoreMarker = insertTryCatchBlocksMarkers(methodNode)
+        isThereAnyTryCatch = saveStackMarkerForRestoreMarker.isNotEmpty()
+        for ((restore, save) in saveStackMarkerForRestoreMarker) {
+            restoreStackMarkersForSaveMarker.getOrPut(save) { SmartList() }.add(restore)
+        }
 
         val inlineMarkersStack = Stack<AbstractInsnNode>()
 
@@ -57,10 +57,6 @@ internal class FixStackContext(val methodNode: MethodNode) {
                     visitFakeAlwaysTrueIfeq(insnNode)
                 pseudoInsn == PseudoInsn.FAKE_ALWAYS_FALSE_IFEQ ->
                     visitFakeAlwaysFalseIfeq(insnNode)
-                pseudoInsn == PseudoInsn.SAVE_STACK_BEFORE_TRY ->
-                    visitSaveStackBeforeTry(insnNode)
-                pseudoInsn == PseudoInsn.RESTORE_STACK_IN_TRY_CATCH ->
-                    visitRestoreStackInTryCatch(insnNode)
                 InlineCodegenUtil.isBeforeInlineMarker(insnNode) -> {
                     inlineMarkersStack.push(insnNode)
                 }
@@ -92,51 +88,18 @@ internal class FixStackContext(val methodNode: MethodNode) {
         fakeAlwaysFalseIfeqMarkers.add(insnNode)
     }
 
-    private fun visitSaveStackBeforeTry(insnNode: AbstractInsnNode) {
-        val tryStartLabel = insnNode.next
-        assert(tryStartLabel is LabelNode) { "${indexOf(insnNode)}: save should be followed by a label" }
-        saveStackNodesForTryStartLabel[tryStartLabel as LabelNode] = insnNode
-    }
-
-    private fun visitRestoreStackInTryCatch(insnNode: AbstractInsnNode) {
-        val restoreLabel = insnNode.findPreviousOrNull { it.hasOpcode() }!!.findPreviousOrNull { it is LabelNode || it.hasOpcode() }!!
-        if (restoreLabel !is LabelNode) {
-            throw AssertionError("${indexOf(insnNode)}: restore should be preceded by a catch block label")
-        }
-        val saveNodes = findMatchingSaveNodes(restoreLabel)
-        if (saveNodes.isEmpty()) {
-            throw AssertionError("${indexOf(insnNode)}: in handler ${indexOf(restoreLabel)} restore is not matched with save")
-        }
-        else if (saveNodes.size > 1) {
-            throw AssertionError("${indexOf(insnNode)}: in handler ${indexOf(restoreLabel)} restore is matched with several saves")
-        }
-        val saveNode = saveNodes.first()
-        saveStackMarkerForRestoreMarker[insnNode] = saveNode
-        restoreStackMarkersForSaveMarker.getOrPut(saveNode, { SmartList<AbstractInsnNode>() }).add(insnNode)
-    }
-
-    private fun findMatchingSaveNodes(restoreLabel: LabelNode): List<AbstractInsnNode> {
-        val saveNodes = SmartHashSet<AbstractInsnNode>()
-        methodNode.tryCatchBlocks.forEach { tcb ->
-            if (restoreLabel == tcb.start || restoreLabel == tcb.handler) {
-                saveStackNodesForTryStartLabel[tcb.start]?.let { saveNodes.add(it) }
-            }
-        }
-        return SmartList<AbstractInsnNode>(saveNodes)
-    }
-
     private fun indexOf(node: AbstractInsnNode) = methodNode.instructions.indexOf(node)
 
     fun hasAnyMarkers(): Boolean =
             breakContinueGotoNodes.isNotEmpty() ||
             fakeAlwaysTrueIfeqMarkers.isNotEmpty() ||
             fakeAlwaysFalseIfeqMarkers.isNotEmpty() ||
-            saveStackNodesForTryStartLabel.isNotEmpty() ||
+            isThereAnyTryCatch ||
             openingInlineMethodMarker.isNotEmpty()
 
     fun isAnalysisRequired(): Boolean =
             breakContinueGotoNodes.isNotEmpty() ||
-            saveStackNodesForTryStartLabel.isNotEmpty() ||
+            isThereAnyTryCatch ||
             openingInlineMethodMarker.isNotEmpty()
 
 }