Proper resort variables on inlining lowered ir closures

Original problem is that lowered ir closures doesn't meet inliner expectations
 about captured variable position in inlining method.
 E.g.: Call 'foo(valueParam) { capturedParam }' to
  inline function 'foo' with declaration

      inline fun foo(valueParam: Foo, inlineParamWithCaptured: Bar.() ->) ....

 is reorganized through inlining to equivalent call foo(valueParam, capturedParam1, cp2 ...).
 But lowered closure for lambda parameter has totally different parameters order:

     fun loweredLambda$x(extensionReceiver, captured1, cp2..., valueParam1, vp2...)

 So before inlining lowered closure should be transformed to

     fun loweredLambda$x(extensionReceiver, valueParam1, vp2..., captured1, cp2..)

 #KT-28547 Fixed

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

@@ -418,12 +418,15 @@ class MethodInliner(
             private val isInliningLambda = nodeRemapper.isInsideInliningLambda
 
             private fun getNewIndex(`var`: Int): Int {
-                if (inliningContext.isInliningLambda && inliningContext.lambdaInfo is IrExpressionLambda) {
+                val lambdaInfo = inliningContext.lambdaInfo
+                if (inliningContext.isInliningLambda && lambdaInfo is IrExpressionLambda) {
                     if (`var` < parameters.argsSizeOnStack) {
-                        if (`var` < capturedParamsSize) {
+                        val capturedParamsStartIndex =
-                            return `var` + realParametersSize
+                            if (lambdaInfo.isExtensionLambda) lambdaInfo.invokeMethod.argumentTypes[0].size else 0 //shift by extension
-                        }
+                        val capturedParamsEndIndex = capturedParamsSize + capturedParamsStartIndex - 1
-                        else {
+                        if (`var` in capturedParamsStartIndex..capturedParamsEndIndex) {
+                            return `var` + realParametersSize - capturedParamsStartIndex //subtract extension
+                        } else if (`var` >= capturedParamsStartIndex) {
                             return `var` - capturedParamsSize
                         }
                     }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.backend.jvm.codegen
 
+import org.jetbrains.kotlin.builtins.isExtensionFunctionType
 import org.jetbrains.kotlin.codegen.Callable
 import org.jetbrains.kotlin.codegen.JvmKotlinType
 import org.jetbrains.kotlin.codegen.StackValue
@@ -34,9 +35,8 @@ class IrInlineCodegen(
         val lambdaInfo = next as IrExpressionLambda
         activeLambda = lambdaInfo
 
-        val argumentTypes = lambdaInfo.loweredMethod.argumentTypes
         lambdaInfo.reference.getArguments().forEachIndexed { index, (_, ir) ->
-            putCapturedValueOnStack(ir, argumentTypes[index], index)
+            putCapturedValueOnStack(ir, lambdaInfo.capturedParamsInDesc[index], index)
         }
         activeLambda = null
     }
@@ -70,10 +70,10 @@ class IrInlineCodegen(
         putArgumentOrCapturedToLocalVal(JvmKotlinType(value.type, value.kotlinType), value, -1, parameterIndex, ValueKind.CAPTURED /*kind*/)
     }
 
-    private fun putCapturedValueOnStack(argumentExpression: IrExpression, valueType: Type, capturedParamindex: Int) {
+    private fun putCapturedValueOnStack(argumentExpression: IrExpression, valueType: Type, capturedParamIndex: Int) {
         val onStack = codegen.gen(argumentExpression, valueType, BlockInfo.create())
         putArgumentOrCapturedToLocalVal(
-            JvmKotlinType(onStack.type, onStack.kotlinType), onStack, capturedParamindex, capturedParamindex, ValueKind.CAPTURED
+            JvmKotlinType(onStack.type, onStack.kotlinType), onStack, capturedParamIndex, capturedParamIndex, ValueKind.CAPTURED
         )
     }
 
@@ -93,10 +93,10 @@ class IrInlineCodegen(
 
     private fun rememberClosure(irReference: IrFunctionReference, type: Type, parameter: ValueParameterDescriptor): LambdaInfo {
         //assert(InlineUtil.isInlinableParameterExpression(ktLambda)) { "Couldn't find inline expression in ${expression.text}" }
-
         val expression = irReference.symbol.owner as IrFunction
         return IrExpressionLambda(
-            irReference, expression, typeMapper, parameter.isCrossinline, false/*TODO*/
+            irReference, expression, typeMapper, parameter.isCrossinline, false/*TODO*/,
+            parameter.type.isExtensionFunctionType
         ).also { lambda ->
             val closureInfo = invocationParamBuilder.addNextValueParameter(type, true, null, parameter.index)
             closureInfo.lambda = lambda
@@ -110,7 +110,8 @@ class IrExpressionLambda(
     val function: IrFunction,
     typeMapper: KotlinTypeMapper,
     isCrossInline: Boolean,
-    override val isBoundCallableReference: Boolean
+    override val isBoundCallableReference: Boolean,
+    val isExtensionLambda: Boolean
 ) : ExpressionLambda(typeMapper, isCrossInline) {
 
     override fun isMyLabel(name: String): Boolean {
@@ -118,30 +119,35 @@ class IrExpressionLambda(
         return false
     }
 
-    override val lambdaClassType: Type
+    override val lambdaClassType: Type = Type.getObjectType("test123")
-        get() = Type.getObjectType("test123")
 
-    override val capturedVars: List<CapturedParamDesc> by lazy {
+    override val capturedVars: List<CapturedParamDesc> =
         arrayListOf<CapturedParamDesc>().apply {
             reference.getArguments().forEachIndexed { _, (_, ir) ->
                 val getValue = ir as? IrGetValue ?: error("Unrecognized expression: $ir")
                 add(capturedParamDesc(getValue.descriptor.name.asString(), typeMapper.mapType(getValue.descriptor.type)))
             }
         }
-    }
 
-    val loweredMethod: Method
+    private val loweredMethod = typeMapper.mapAsmMethod(function.descriptor)
-        get() = typeMapper.mapAsmMethod(function.descriptor)
+
+    val capturedParamsInDesc: List<Type> =
+        loweredMethod.argumentTypes.drop(if (isExtensionLambda) 1 else 0).take(capturedVars.size)
 
     override val invokeMethod: Method = loweredMethod.let {
-        Method(it.name, it.returnType, it.argumentTypes.drop(capturedVars.size).toTypedArray())
+        Method(
+            it.name,
+            it.returnType,
+            (
+                    (if (isExtensionLambda) it.argumentTypes.take(1) else emptyList()) +
+                            it.argumentTypes.drop((if (isExtensionLambda) 1 else 0) + capturedVars.size)
+                    ).toTypedArray()
+        )
     }
 
-    override val invokeMethodDescriptor: FunctionDescriptor
+    override val invokeMethodDescriptor: FunctionDescriptor = function.descriptor
-        get() = function.descriptor
 
-    override val hasDispatchReceiver: Boolean
+    override val hasDispatchReceiver: Boolean = false
-        get() = false
 }
 
 fun isInlineIrExpression(argumentExpression: IrExpression) =

compiler/testData/codegen/box/defaultArguments/kt6382.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 
 fun box(): String {

compiler/testData/codegen/box/jvmName/classMembers.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 
 // WITH_RUNTIME

compiler/testData/codegen/box/jvmStatic/convention.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 
 // WITH_RUNTIME

compiler/testData/codegen/box/objects/objectVsClassInitialization_kt5291.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 public inline fun <T> T.with(f: T.() -> Unit): T {
     this.f()
     return this

compiler/testData/codegen/box/properties/lateinit/isInitializedAndDeinitialize/sideEffects.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 
 class Foo {

compiler/testData/codegen/box/ranges/forInStringWithCustomIterator.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // WITH_RUNTIME
 

compiler/testData/codegen/boxInline/callableReference/kt15449.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // FILE: 1.kt
 
 package test

compiler/testData/codegen/boxInline/complex/swapAndWith.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // NO_CHECK_LAMBDA_INLINING
 // FILE: 1.kt
 package test

compiler/testData/codegen/boxInline/complex/swapAndWith2.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // FILE: 1.kt
 package test
 

compiler/testData/codegen/boxInline/localFunInLambda/defaultParam.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // FILE: 1.kt
 
 inline fun <T> run(c: () -> T): T = c()

compiler/testData/codegen/boxInline/localFunInLambda/localFunInLambda.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // FILE: 1.kt
 // WITH_RUNTIME
 package test

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

@@ -0,0 +1,52 @@
+// FILE: 1.kt
+package test
+
+class C {
+    var inserting: Boolean = false
+    fun nextSlot(): Any? = null
+    fun startNode(key: Any?) {}
+    fun endNode() {}
+    fun emitNode(node: Any?) {}
+    fun useNode(): Any? = null
+    fun skipValue() {}
+    fun updateValue(value: Any?) {}
+}
+
+class B<T>(val composer: C, val node: T) {
+    inline fun <V> bar(value: V, block: T.(V) -> Unit) = with(composer) {
+        if (inserting || nextSlot() != value) {
+            updateValue(value)
+            node.block(value)
+        } else skipValue()
+    }
+}
+
+class A(val composer: C) {
+    inline fun <T> foo(key: Any, ctor: () -> T, update: B<T>.() -> Unit) = with(composer) {
+        startNode(key)
+        val node = if (inserting)
+            ctor().also { emitNode(it) }
+        else useNode() as T
+        B<T>(this, node).update()
+        endNode()
+    }
+}
+
+// FILE: 2.kt
+import test.*
+
+fun box(): String {
+    val a = A(C())
+    val str = "OK"
+    var result = "fail"
+    a.foo<String>(
+        123,
+        { "abc" },
+        {
+            bar(str) { }
+            result = "OK"
+        }
+    )
+
+    return result
+}

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

@@ -0,0 +1,21 @@
+// FILE: 1.kt
+
+package test
+
+inline fun Double.run(body : Double.() -> String): String {
+    return this.body()
+}
+
+// FILE: 2.kt
+
+import test.*
+
+fun box(): String {
+    var captured = "fail"
+    return 1.0.run {
+        if (this == 1.0) {
+             "OK"
+        } else captured
+    }
+
+}

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

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // FILE: 1.kt
 
 package test

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

@@ -2813,6 +2813,16 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
             runTest("compiler/testData/codegen/boxInline/simple/kt17431.kt");
         }
 
+        @TestMetadata("kt28547.kt")
+        public void testKt28547() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/kt28547.kt");
+        }
+
+        @TestMetadata("kt28547_2.kt")
+        public void testKt28547_2() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/kt28547_2.kt");
+        }
+
         @TestMetadata("params.kt")
         public void testParams() throws Exception {
             runTest("compiler/testData/codegen/boxInline/simple/params.kt");

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

@@ -2813,6 +2813,16 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
             runTest("compiler/testData/codegen/boxInline/simple/kt17431.kt");
         }
 
+        @TestMetadata("kt28547.kt")
+        public void testKt28547() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/kt28547.kt");
+        }
+
+        @TestMetadata("kt28547_2.kt")
+        public void testKt28547_2() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/kt28547_2.kt");
+        }
+
         @TestMetadata("params.kt")
         public void testParams() throws Exception {
             runTest("compiler/testData/codegen/boxInline/simple/params.kt");

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

@@ -2813,6 +2813,16 @@ public class IrBlackBoxInlineCodegenTestGenerated extends AbstractIrBlackBoxInli
             runTest("compiler/testData/codegen/boxInline/simple/kt17431.kt");
         }
 
+        @TestMetadata("kt28547.kt")
+        public void testKt28547() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/kt28547.kt");
+        }
+
+        @TestMetadata("kt28547_2.kt")
+        public void testKt28547_2() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/simple/kt28547_2.kt");
+        }
+
         @TestMetadata("params.kt")
         public void testParams() throws Exception {
             runTest("compiler/testData/codegen/boxInline/simple/params.kt");