JVM_IR KT-47984 use stack size calculator from ASM

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

@@ -7,11 +7,11 @@ package org.jetbrains.kotlin.codegen.inline
 
 import org.jetbrains.kotlin.codegen.optimization.boxing.isMethodInsnWith
 import org.jetbrains.kotlin.codegen.optimization.common.InsnSequence
-import org.jetbrains.kotlin.codegen.optimization.common.remapLocalVariables
+import org.jetbrains.kotlin.codegen.optimization.common.removeUnusedLocalVariables
+import org.jetbrains.kotlin.codegen.optimization.common.updateMaxStack
 import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
 import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Opcodes
-import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.tree.*
 
 class InplaceArgumentsMethodTransformer : MethodTransformer() {
@@ -27,8 +27,8 @@ class InplaceArgumentsMethodTransformer : MethodTransformer() {
             transformMethod(methodContext)
             updateLvtEntriesForMovedInstructions(methodContext)
 
-            methodNode.maxStack = StackSizeCalculator(internalClassName, methodNode).calculateStackSize()
+            methodNode.removeUnusedLocalVariables()
-            packLocalVariables(methodNode)
+            methodNode.updateMaxStack()
         }
         stripMarkers(methodNode)
     }
@@ -383,61 +383,4 @@ class InplaceArgumentsMethodTransformer : MethodTransformer() {
             insn = insn.next
         }
     }
-
-    private fun packLocalVariables(methodNode: MethodNode) {
-        // After we've dropped argument stores-loads for inplace arguments, some variable slots may become unused.
-        // Track used slots and remap local variables.
-        // Keep in mind that 'long' and 'double' variables occupy 2 slots.
-        val usedLocalVar = BooleanArray(methodNode.maxLocals)
-
-        // 'this' and arguments are always "used"
-        val argTypes = Type.getArgumentTypes(methodNode.desc)
-        var lastArgIndex = 0
-        if (methodNode.access and Opcodes.ACC_STATIC == 0) {
-            usedLocalVar[lastArgIndex++] = true
-        }
-        for (argType in argTypes) {
-            usedLocalVar[lastArgIndex++] = true
-            if (argType.size == 2) {
-                usedLocalVar[lastArgIndex++] = true
-            }
-        }
-
-        // Local variables used in xLOAD/xSTORE/IINC instructions
-        for (insn in methodNode.instructions) {
-            when (insn.opcode) {
-                Opcodes.ILOAD, Opcodes.FLOAD, Opcodes.ALOAD, Opcodes.ISTORE, Opcodes.FSTORE, Opcodes.ASTORE ->
-                    usedLocalVar[(insn as VarInsnNode).`var`] = true
-                Opcodes.LLOAD, Opcodes.DLOAD, Opcodes.LSTORE, Opcodes.DSTORE -> {
-                    val index = (insn as VarInsnNode).`var`
-                    usedLocalVar[index] = true
-                    usedLocalVar[index + 1] = true
-                }
-                Opcodes.IINC ->
-                    usedLocalVar[(insn as IincInsnNode).`var`] = true
-            }
-        }
-
-        // Local variables mentioned in LVT
-        for (lv in methodNode.localVariables) {
-            usedLocalVar[lv.index] = true
-            val lvd0 = lv.desc[0]
-            if (lvd0 == 'J' || lvd0 == 'D') { // long || double
-                usedLocalVar[lv.index + 1] = true
-            }
-        }
-
-        val newIndex = IntArray(methodNode.maxLocals)
-        var lastIndex = 0
-        for (i in newIndex.indices) {
-            newIndex[i] = lastIndex
-            if (usedLocalVar[i]) {
-                ++lastIndex
-            }
-        }
-
-        methodNode.remapLocalVariables(newIndex)
-        methodNode.maxLocals = lastIndex
-    }
-
 }

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

@@ -1,79 +0,0 @@
-/*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.codegen.inline
-
-import com.intellij.util.containers.Stack
-import org.jetbrains.kotlin.codegen.optimization.fixStack.FastStackAnalyzer
-import org.jetbrains.kotlin.codegen.optimization.fixStack.FixStackInterpreter
-import org.jetbrains.kotlin.codegen.optimization.fixStack.FixStackValue
-import org.jetbrains.org.objectweb.asm.Opcodes
-import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
-import org.jetbrains.org.objectweb.asm.tree.MethodNode
-import org.jetbrains.org.objectweb.asm.tree.analysis.Frame
-import org.jetbrains.org.objectweb.asm.tree.analysis.Interpreter
-
-internal class StackSizeCalculator(owner: String, method: MethodNode) :
-    FastStackAnalyzer<FixStackValue>(owner, method, FixStackInterpreter()) {
-
-    fun calculateStackSize(): Int {
-        val frames = analyze()
-        return frames.maxOf { frame ->
-            if (frame is ExpandableStackFrame) frame.getActualStackSize() else method.maxStack
-        }
-    }
-
-    override fun newFrame(nLocals: Int, nStack: Int): Frame<FixStackValue> =
-        ExpandableStackFrame(nLocals, nStack)
-
-    class ExpandableStackFrame(nLocals: Int, nStack: Int) : Frame<FixStackValue>(nLocals, nStack) {
-        private val extraStack = Stack<FixStackValue>()
-
-        override fun init(src: Frame<out FixStackValue>): Frame<FixStackValue> {
-            extraStack.clear()
-            extraStack.addAll((src as ExpandableStackFrame).extraStack)
-            return super.init(src)
-        }
-
-        override fun clearStack() {
-            extraStack.clear()
-            super.clearStack()
-        }
-
-        override fun execute(insn: AbstractInsnNode, interpreter: Interpreter<FixStackValue>) {
-            if (insn.opcode == Opcodes.RETURN) return
-            super.execute(insn, interpreter)
-        }
-
-        override fun push(value: FixStackValue) {
-            if (super.getStackSize() < maxStackSize) {
-                super.push(value)
-            } else {
-                extraStack.add(value)
-            }
-        }
-
-        override fun pop(): FixStackValue =
-            if (extraStack.isNotEmpty()) {
-                extraStack.pop()
-            } else {
-                super.pop()
-            }
-
-        override fun setStack(i: Int, value: FixStackValue) {
-            if (i < maxStackSize) {
-                super.setStack(i, value)
-            } else {
-                extraStack[i - maxStackSize] = value
-            }
-        }
-
-        override fun merge(frame: Frame<out FixStackValue>, interpreter: Interpreter<FixStackValue>): Boolean {
-            throw UnsupportedOperationException("Stack size calculator should not merge frames")
-        }
-
-        fun getActualStackSize() = super.getStackSize() + extraStack.sumOf { it.size }
-    }
-}

compiler/backend/src/org/jetbrains/kotlin/codegen/optimization/common/Util.kt

@@ -22,7 +22,6 @@ import org.jetbrains.kotlin.codegen.optimization.removeNodeGetNext
 import org.jetbrains.kotlin.codegen.pseudoInsns.PseudoInsn
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.org.objectweb.asm.MethodVisitor
-import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Opcodes.*
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.tree.*
@@ -36,7 +35,7 @@ val AbstractInsnNode.isMeaningful: Boolean
 
 val AbstractInsnNode.isBranchOrCall: Boolean
     get() =
-        when(this.type) {
+        when (this.type) {
             AbstractInsnNode.JUMP_INSN,
             AbstractInsnNode.TABLESWITCH_INSN,
             AbstractInsnNode.LOOKUPSWITCH_INSN,
@@ -85,13 +84,17 @@ fun MethodNode.prepareForEmitting() {
 
         current = prev
     }
+    updateMaxStack()
+}
+
+fun MethodNode.updateMaxStack() {
     maxStack = -1
     accept(
         MaxStackFrameSizeAndLocalsCalculator(
-            Opcodes.API_VERSION, access, desc,
+            API_VERSION, access, desc,
-            object : MethodVisitor(Opcodes.API_VERSION) {
+            object : MethodVisitor(API_VERSION) {
                 override fun visitMaxs(maxStack: Int, maxLocals: Int) {
-                    this@prepareForEmitting.maxStack = maxStack
+                    this@updateMaxStack.maxStack = maxStack
                 }
             })
     )
@@ -100,10 +103,10 @@ fun MethodNode.prepareForEmitting() {
 fun MethodNode.stripOptimizationMarkers() {
     var insn = instructions.first
     while (insn != null) {
-        if (isOptimizationMarker(insn)) {
+        insn = if (isOptimizationMarker(insn)) {
-            insn = instructions.removeNodeGetNext(insn)
+            instructions.removeNodeGetNext(insn)
         } else {
-            insn = insn.next
+            insn.next
         }
     }
 }
@@ -123,7 +126,7 @@ fun MethodNode.removeUnusedLocalVariables() {
     // Arguments are always used whether or not they are in the local variable table
     // or used by instructions.
     var argumentIndex = 0
-    val isStatic = (access and Opcodes.ACC_STATIC) != 0
+    val isStatic = (access and ACC_STATIC) != 0
     if (!isStatic) {
         used[argumentIndex++] = true
     }
@@ -230,8 +233,8 @@ val AbstractInsnNode.intConstant: Int?
 
 fun insnListOf(vararg insns: AbstractInsnNode) = InsnList().apply { insns.forEach { add(it) } }
 
-fun AbstractInsnNode.isStoreOperation(): Boolean = opcode in Opcodes.ISTORE..Opcodes.ASTORE
+fun AbstractInsnNode.isStoreOperation(): Boolean = opcode in ISTORE..ASTORE
-fun AbstractInsnNode.isLoadOperation(): Boolean = opcode in Opcodes.ILOAD..Opcodes.ALOAD
+fun AbstractInsnNode.isLoadOperation(): Boolean = opcode in ILOAD..ALOAD
 
 val AbstractInsnNode?.debugText
     get() =

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -18275,6 +18275,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
         public void testSuspensionPointInsideArgument() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
         }
+
+        @Test
+        @TestMetadata("withLogFile.kt")
+        public void testWithLogFile() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/withLogFile.kt");
+        }
     }
 
     @Nested

compiler/testData/codegen/box/inlineArgsInPlace/withLogFile.kt

@@ -0,0 +1,35 @@
+// TARGET_BACKEND: JVM
+// WITH_STDLIB
+// FULL_JDK
+
+fun condition1() = true
+
+fun zzz() {}
+
+fun f2() = 2
+
+// Minimized version of
+// https://github.com/JetBrains/kotlin/commit/ced973b7074f4207859d9709375f2bf28b3e2c55#diff-f9a8dce85985573b5478da1b5379342fe37fca94c14f55b69d1c884fece42f92R841
+fun box(): String {
+    val arr = arrayOfNulls<Int>(4)
+
+    fun zap(threadNo: Int): String {
+        arr[threadNo] = try {
+            f2()
+        } catch (e: Exception) {
+            null
+        }
+        arr[threadNo] = when {
+            condition1() -> {
+                1
+            }
+            else -> {
+                println("[$threadNo] 3")
+                3
+            }
+        }
+        return "OK"
+    }
+
+    return zap(0)
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -18131,6 +18131,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
         public void testSuspensionPointInsideArgument() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
         }
+
+        @Test
+        @TestMetadata("withLogFile.kt")
+        public void testWithLogFile() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/withLogFile.kt");
+        }
     }
 
     @Nested

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -18275,6 +18275,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
         public void testSuspensionPointInsideArgument() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
         }
+
+        @Test
+        @TestMetadata("withLogFile.kt")
+        public void testWithLogFile() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/withLogFile.kt");
+        }
     }
 
     @Nested

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

@@ -15028,6 +15028,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
         public void testSuspensionPointInsideArgument() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
         }
+
+        @TestMetadata("withLogFile.kt")
+        public void testWithLogFile() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/withLogFile.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/inlineClasses")