Handle DUPnXm instructions in PopBackwardPropagationTransformer

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/boxing/PopBackwardPropagationTransformer.kt

@@ -17,8 +17,10 @@
 package org.jetbrains.kotlin.codegen.optimization.boxing
 
 import org.jetbrains.kotlin.codegen.optimization.OptimizationMethodVisitor
+import org.jetbrains.kotlin.codegen.optimization.common.debugText
 import org.jetbrains.kotlin.codegen.optimization.common.isLoadOperation
 import org.jetbrains.kotlin.codegen.optimization.common.isMeaningful
+import org.jetbrains.kotlin.codegen.optimization.fixStack.peekWords
 import org.jetbrains.kotlin.codegen.optimization.fixStack.top
 import org.jetbrains.kotlin.codegen.optimization.removeNodeGetNext
 import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
@@ -33,17 +35,9 @@ import java.util.*
 class PopBackwardPropagationTransformer : MethodTransformer() {
     override fun transform(internalClassName: String, methodNode: MethodNode) {
         if (!OptimizationMethodVisitor.canBeOptimizedUsingSourceInterpreter(methodNode)) return
-        if (methodNode.instructions.toArray().any { it.isUnsafeStackInsn() }) return
         Transformer(methodNode).transform()
     }
 
-    // TODO better stack operations analysis
-    private fun AbstractInsnNode.isUnsafeStackInsn() =
-            opcode == Opcodes.DUP_X1 ||
-            opcode == Opcodes.DUP_X2 ||
-            opcode == Opcodes.DUP2_X1 ||
-            opcode == Opcodes.DUP2_X2
-
     private class Transformer(val methodNode: MethodNode) {
         private interface Transformation {
             fun apply(insn: AbstractInsnNode)
@@ -82,8 +76,46 @@ class PopBackwardPropagationTransformer : MethodTransformer() {
             postprocessNops()
         }
 
-        private fun analyzeMethodBody(): Array<out Frame<SourceValue>?> =
+        private fun analyzeMethodBody(): Array<out Frame<SourceValue>?> {
-                Analyzer<SourceValue>(object : SourceInterpreter() {
+            val frames = Analyzer<SourceValue>(HazardsTrackingInterpreter()).analyze("fake", methodNode)
+
+            postprocessDupNxM(frames)
+
+            return frames
+        }
+
+        private fun postprocessDupNxM(frames: Array<out Frame<SourceValue>?>) {
+            val insns = methodNode.instructions.toArray()
+            for (i in frames.indices) {
+                val frame = frames[i] ?: continue
+                val insn = insns[i]
+
+                when (insn.opcode) {
+                    Opcodes.DUP_X1 -> {
+                        val top2 = frame.peekWords(1, 1) ?: throwIncorrectBytecode(insn, frame)
+                        top2.forEach { it.insns.markAsDontTouch() }
+                    }
+                    Opcodes.DUP2_X1 -> {
+                        val top3 = frame.peekWords(2, 1) ?: throwIncorrectBytecode(insn, frame)
+                        top3.forEach { it.insns.markAsDontTouch() }
+                    }
+                    Opcodes.DUP_X2 -> {
+                        val top3 = frame.peekWords(1, 2) ?: throwIncorrectBytecode(insn, frame)
+                        top3.forEach { it.insns.markAsDontTouch() }
+                    }
+                    Opcodes.DUP2_X2 -> {
+                        val top4 = frame.peekWords(2, 2) ?: throwIncorrectBytecode(insn, frame)
+                        top4.forEach { it.insns.markAsDontTouch() }
+                    }
+                }
+            }
+        }
+
+        private fun throwIncorrectBytecode(insn: AbstractInsnNode?, frame: Frame<SourceValue>): Nothing {
+            throw AssertionError("Incorrect bytecode at ${methodNode.instructions.indexOf(insn)}: ${insn.debugText} $frame")
+        }
+
+        private inner class HazardsTrackingInterpreter : SourceInterpreter() {
             override fun naryOperation(insn: AbstractInsnNode, values: MutableList<out SourceValue>): SourceValue {
                 for (value in values) {
                     value.insns.markAsDontTouch()
@@ -115,13 +147,13 @@ class PopBackwardPropagationTransformer : MethodTransformer() {
                 value3.insns.markAsDontTouch()
                 return super.ternaryOperation(insn, value1, value2, value3)
             }
+        }
 
         private fun Collection<AbstractInsnNode>.markAsDontTouch() {
             forEach {
                 dontTouchInsnIndices[insnList.indexOf(it)] = true
             }
         }
-                }).analyze("fake", methodNode)
 
 
         private fun computeTransformations() {

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

@@ -28,6 +28,31 @@ fun <V : Value> Frame<V>.top(): V? =
 fun <V : Value> Frame<V>.peek(offset: Int): V? =
         if (stackSize > offset) getStack(stackSize - offset - 1) else null
 
+private fun <V : Value> Frame<V>.peekWordsTo(dest: MutableList<V>, size: Int, offset0: Int = 0): Int {
+    var offset = offset0
+    var totalSize = 0
+    while (totalSize < size) {
+        val value = peek(offset++) ?: return -1
+        dest.add(value)
+        totalSize += value.size
+    }
+    if (totalSize > size) return -1
+    return offset
+}
+
+fun <V : Value> Frame<V>.peekWords(size: Int): List<V>? {
+    val result = ArrayList<V>(size)
+    return if (peekWordsTo(result, size) < 0) null else result
+}
+
+fun <V : Value> Frame<V>.peekWords(size1: Int, size2: Int): List<V>? {
+    val result = ArrayList<V>(size1 + size2)
+    val offset = peekWordsTo(result, size1)
+    if (offset < 0) return null
+    if (peekWordsTo(result, size2, offset) < 0) return null
+    return result
+}
+
 class SavedStackDescriptor(
         val savedValues: List<BasicValue>,
         val firstLocalVarIndex: Int