Support simple capturing in IR inliner

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

@@ -16,9 +16,11 @@
 
 package org.jetbrains.kotlin.codegen.inline
 
+import org.jetbrains.kotlin.backend.jvm.codegen.IrExpressionLambda
 import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.ClosureCodegen
 import org.jetbrains.kotlin.codegen.StackValue
+import org.jetbrains.kotlin.codegen.inline.FieldRemapper.Companion.foldName
 import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethods
 import org.jetbrains.kotlin.codegen.optimization.ApiVersionCallsPreprocessingMethodTransformer
 import org.jetbrains.kotlin.codegen.optimization.FixStackWithLabelNormalizationMethodTransformer
@@ -328,17 +330,29 @@ class MethodInliner(
 
         val capturedParamsSize = parameters.capturedParametersSizeOnStack
         val realParametersSize = parameters.realParametersSizeOnStack
+        val transformedNode = MethodNode(
+            API, node.access, node.name,
+            Type.getMethodDescriptor(Type.getReturnType(node.desc), *(Type.getArgumentTypes(node.desc) + parameters.capturedTypes)),
+            node.signature, node.exceptions?.toTypedArray()
+        )
 
-        val transformedNode = object : MethodNode(
+        val transformationVisitor = object : MethodVisitor(API, transformedNode) {
-                API, node.access, node.name,
-                Type.getMethodDescriptor(Type.getReturnType(node.desc), *(Type.getArgumentTypes(node.desc) + parameters.capturedTypes)),
-                node.signature, node.exceptions?.toTypedArray()
-        ) {
             private val GENERATE_DEBUG_INFO = GENERATE_SMAP && inlineOnlySmapSkipper == null
 
             private val isInliningLambda = nodeRemapper.isInsideInliningLambda
 
             private fun getNewIndex(`var`: Int): Int {
+                if (inliningContext.isInliningLambda && inliningContext.lambdaInfo is IrExpressionLambda) {
+                    if (`var` < parameters.argsSizeOnStack) {
+                        if (`var` < capturedParamsSize) {
+                            return `var` + realParametersSize
+                        }
+                        else {
+                            return `var` - capturedParamsSize
+                        }
+                    }
+                    return `var`
+                }
                 return `var` + if (`var` < realParametersSize) 0 else capturedParamsSize
             }
 
@@ -391,7 +405,7 @@ class MethodInliner(
             }
         }
 
-        node.accept(transformedNode)
+        node.accept(transformationVisitor)
 
         transformCaptured(transformedNode)
         transformFinallyDeepIndex(transformedNode, finallyDeepShift)
@@ -643,6 +657,33 @@ class MethodInliner(
             return
         }
 
+        if (inliningContext.isInliningLambda && inliningContext.lambdaInfo is IrExpressionLambda) {
+            val capturedVars = inliningContext.lambdaInfo!!.capturedVars
+            var offset = parameters.realParametersSizeOnStack
+            val map = capturedVars.map {
+                offset to it.also { offset += it.type.size }
+            }.toMap()
+
+            var cur: AbstractInsnNode? = node.instructions.first
+            while (cur != null) {
+                if (cur is VarInsnNode && cur.opcode == Opcodes.ALOAD && map.contains(cur.`var`)) {
+                    val varIndex = cur.`var`
+                    val capturedParamDesc = map[varIndex]!!
+
+                    val newIns = FieldInsnNode(
+                        Opcodes.GETSTATIC,
+                        capturedParamDesc.containingLambdaName,
+                        foldName(capturedParamDesc.fieldName),
+                        capturedParamDesc.type.descriptor
+                    )
+                    node.instructions.insertBefore(cur, newIns)
+                    node.instructions.remove(cur)
+                    cur = newIns
+                }
+                cur = cur.next
+            }
+        }
+
         // Fold all captured variables access chains
         //          ALOAD 0
         //          [ALOAD this$0]*

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

@@ -42,11 +42,13 @@ class IrInlineCodegen(
 
     override fun putClosureParametersOnStack(next: LambdaInfo, functionReferenceReceiver: StackValue?) {
         val lambdaInfo = next as IrExpressionLambda
+        activeLambda = lambdaInfo
 
         val argumentTypes = lambdaInfo.loweredMethod.argumentTypes
-        lambdaInfo.reference.getArguments().forEachIndexed { index, (descriptor, ir) ->
+        lambdaInfo.reference.getArguments().forEachIndexed { index, (_, ir) ->
             putCapturedValueOnStack(ir, argumentTypes[index], index)
         }
+        activeLambda = null
     }
 
     override fun genValueAndPut(valueParameterDescriptor: ValueParameterDescriptor?, argumentExpression: IrExpression, parameterType: Type, parameterIndex: Int, codegen: ExpressionCodegen, blockInfo: BlockInfo) {
@@ -112,18 +114,25 @@ class IrExpressionLambda(
     override val lambdaClassType: Type
         get() = Type.getObjectType("test123")
 
-    override val invokeMethod: Method
+    override val capturedVars: List<CapturedParamDesc> by lazy {
-        get() = typeMapper.mapAsmMethod(function.descriptor)
+        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
         get() = typeMapper.mapAsmMethod(function.descriptor)
 
+    override val invokeMethod: Method = loweredMethod.let {
+        Method(it.name, it.returnType, it.argumentTypes.drop(capturedVars.size).toTypedArray())
+    }
+
     override val invokeMethodDescriptor: FunctionDescriptor
         get() = function.descriptor
 
-    override val capturedVars: List<CapturedParamDesc>
-        get() = emptyList() //cause closure conversion
-
     override val hasDispatchReceiver: Boolean
         get() = false
 }