JVM_IR: move lambda captures to end of signature when inlining

For example, a lambda `{ param -> captured }` of type `E.(T) -> U` will
be transformed by LocalDeclarationsLowering into a private static method

    fun f$lambda-0($this: E, $captured: U, param: T) = $captured

The reason for such an ordering is that a lambda looks the same as a
local function, and local function can have default arguments, and those
arguments can reference captured variables; thus, captured variables
must come before actual declared arguments.

However, this is not the order that the inliner wants. Moreover, since
it was written to handle lambdas represented as `invoke` methods of
anonymous objects, it does not expect the actual callable method to have
any parameters corresponding to captured variables at all. This results
in it attempting to generate a temporary node with descriptor

    (LE;LU;LT;LU;)LU;

while still using locals 1 and 2 as `param` and `$captured` respectively.
In the example above, this is not critical, as they both have reference
type and the lambda will eventually be pasted into a different node
anyway; however, if it happens that one of them is a primitive, or both
are primitives of different types, the bytecode will use incorrect
instructions, causing verification errors. The correct descriptor is

    (LE;LT;LU;)LU;

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

@@ -413,9 +413,19 @@ class MethodInliner(
 
         val capturedParamsSize = parameters.capturedParametersSizeOnStack
         val realParametersSize = parameters.realParametersSizeOnStack
+        val reorderIrLambdaParameters = inliningContext.isInliningLambda &&
+                inliningContext.parent?.isInliningLambda == false &&
+                inliningContext.lambdaInfo is IrExpressionLambda
+        val newArgumentList = if (reorderIrLambdaParameters) {
+            // In IR lambdas, captured variables come before real parameters, but after the extension receiver.
+            // Move them to the end of the descriptor instead.
+            Type.getArgumentTypes(inliningContext.lambdaInfo!!.invokeMethod.descriptor) + parameters.capturedTypes
+        } else {
+            Type.getArgumentTypes(node.desc) + parameters.capturedTypes
+        }
         val transformedNode = MethodNode(
             Opcodes.API_VERSION, node.access, node.name,
-            Type.getMethodDescriptor(Type.getReturnType(node.desc), *(Type.getArgumentTypes(node.desc) + parameters.capturedTypes)),
+            Type.getMethodDescriptor(Type.getReturnType(node.desc), *newArgumentList),
             node.signature, node.exceptions?.toTypedArray()
         )
 
@@ -426,19 +436,21 @@ class MethodInliner(
 
             private fun getNewIndex(`var`: Int): Int {
                 val lambdaInfo = inliningContext.lambdaInfo
-                if (inliningContext.isInliningLambda && lambdaInfo is IrExpressionLambda) {
-                    if (`var` < parameters.argsSizeOnStack) {
-                        val capturedParamsStartIndex =
-                            if (lambdaInfo.isExtensionLambda) lambdaInfo.invokeMethod.argumentTypes[0].size else 0 //shift by extension
-                        val capturedParamsEndIndex = capturedParamsSize + capturedParamsStartIndex - 1
-                        if (`var` in capturedParamsStartIndex..capturedParamsEndIndex) {
-                            return `var` + realParametersSize - capturedParamsStartIndex //subtract extension
-                        } else if (`var` >= capturedParamsStartIndex) {
-                            return `var` - capturedParamsSize
-                        }
+                if (reorderIrLambdaParameters && lambdaInfo is IrExpressionLambda) {
+                    val extensionSize = if (lambdaInfo.isExtensionLambda) lambdaInfo.invokeMethod.argumentTypes[0].size else 0
+                    return when {
+                        //                v-- extensionSize     v-- argsSizeOnStack
+                        // |- extension -|- captured -|- real -|- locals -|    old descriptor
+                        // |- extension -|- real -|- captured -|- locals -|    new descriptor
+                        //                         ^-- realParametersSize
+                        `var` >= parameters.argsSizeOnStack -> `var`
+                        `var` >= extensionSize + capturedParamsSize -> `var` - capturedParamsSize
+                        `var` >= extensionSize -> `var` + realParametersSize - extensionSize
+                        else -> `var`
                     }
-                    return `var`
                 }
+                // |- extension -|- real -|- locals -|               old descriptor
+                // |- extension -|- real -|- captured -|- locals -|  new descriptor
                 return `var` + if (`var` < realParametersSize) 0 else capturedParamsSize
             }
 

compiler/testData/codegen/box/unsignedTypes/unsignedRangeIterator.kt

@@ -1,6 +1,5 @@
 // KJS_WITH_FULL_RUNTIME
 // WITH_RUNTIME
-// IGNORE_BACKEND: JVM_IR
 
 fun testUIntRangeForEach() {
     var s = 0

compiler/testData/codegen/boxInline/simple/captureAndArgumentIncompatibleTypes.kt

@@ -0,0 +1,13 @@
+// FILE: 1.kt
+package test
+
+inline fun foo(i: Int) = i.toFloat()
+
+// FILE: 2.kt
+import test.*
+
+fun box(): String {
+    val captured = 1.0f
+    val result = 1.let { captured + foo(it) }
+    return if (result == 2.0f) "OK" else "Fail: $result"
+}

compiler/testData/codegen/bytecodeText/boxingOptimization/unsignedRangeIteratorSpecialization.kt

@@ -1,5 +1,4 @@
 // WITH_RUNTIME
-// IGNORE_BACKEND: JVM_IR
 
 fun testUIntRangeForEach() {
     var s = 0

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

@@ -3017,6 +3017,11 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/simple"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
         }
 
+        @TestMetadata("captureAndArgumentIncompatibleTypes.kt")
+        public void testCaptureAndArgumentIncompatibleTypes() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/captureAndArgumentIncompatibleTypes.kt");
+        }
+
         @TestMetadata("classObject.kt")
         public void testClassObject() throws Exception {
             runTest("compiler/testData/codegen/boxInline/simple/classObject.kt");

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

@@ -3017,6 +3017,11 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/simple"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
         }
 
+        @TestMetadata("captureAndArgumentIncompatibleTypes.kt")
+        public void testCaptureAndArgumentIncompatibleTypes() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/captureAndArgumentIncompatibleTypes.kt");
+        }
+
         @TestMetadata("classObject.kt")
         public void testClassObject() throws Exception {
             runTest("compiler/testData/codegen/boxInline/simple/classObject.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxInlineCodegenTestGenerated.java

@@ -3017,6 +3017,11 @@ public class IrBlackBoxInlineCodegenTestGenerated extends AbstractIrBlackBoxInli
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/simple"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM_IR, true);
         }
 
+        @TestMetadata("captureAndArgumentIncompatibleTypes.kt")
+        public void testCaptureAndArgumentIncompatibleTypes() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/captureAndArgumentIncompatibleTypes.kt");
+        }
+
         @TestMetadata("classObject.kt")
         public void testClassObject() throws Exception {
             runTest("compiler/testData/codegen/boxInline/simple/classObject.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrCompileKotlinAgainstInlineKotlinTestGenerated.java

@@ -3017,6 +3017,11 @@ public class IrCompileKotlinAgainstInlineKotlinTestGenerated extends AbstractIrC
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/simple"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM_IR, true);
         }
 
+        @TestMetadata("captureAndArgumentIncompatibleTypes.kt")
+        public void testCaptureAndArgumentIncompatibleTypes() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/captureAndArgumentIncompatibleTypes.kt");
+        }
+
         @TestMetadata("classObject.kt")
         public void testClassObject() throws Exception {
             runTest("compiler/testData/codegen/boxInline/simple/classObject.kt");