JVM: keep track of which regenerated field is for which parameter

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

@@ -129,13 +129,13 @@ class AnonymousObjectTransformer(
 
         val allCapturedParamBuilder = ParametersBuilder.newBuilder()
         val constructorParamBuilder = ParametersBuilder.newBuilder()
-        val additionalFakeParams = extractParametersMappingAndPatchConstructor(
+        extractParametersMappingAndPatchConstructor(
             constructor!!, allCapturedParamBuilder, constructorParamBuilder, transformationInfo, parentRemapper
         )
 
         val deferringMethods = ArrayList<DeferredMethodVisitor>()
 
-        generateConstructorAndFields(classBuilder, allCapturedParamBuilder, constructorParamBuilder, parentRemapper, additionalFakeParams)
+        generateConstructorAndFields(classBuilder, constructorParamBuilder, parentRemapper)
 
         val coroutineTransformer = CoroutineTransformer(
             inliningContext,
@@ -341,10 +341,8 @@ class AnonymousObjectTransformer(
 
     private fun generateConstructorAndFields(
         classBuilder: ClassBuilder,
-        allCapturedBuilder: ParametersBuilder,
         constructorInlineBuilder: ParametersBuilder,
-        parentRemapper: FieldRemapper,
-        constructorAdditionalFakeParams: List<CapturedParamInfo>
+        parentRemapper: FieldRemapper
     ) {
         val constructorParams = constructorInlineBuilder.buildParameters()
         val constructorParamTypes = constructorParams.filter { !it.isSkipped }.map { it.type }.drop(1)
@@ -359,47 +357,24 @@ class AnonymousObjectTransformer(
         constructorVisitor.visitLabel(newBodyStartLabel)
 
         //initialize captured fields
-        val capturedIndexes = IntArray(constructorParams.parameters.size)
-        var index = 0
-        var size = 0
-        for (info in constructorParams) {
-            if (!info.isSkipped) { //not inlined
-                if (info.isCaptured || info is CapturedParamInfo) {
-                    capturedIndexes[index] = size
-                    index++
-                }
-                size += info.type.size
-            }
-        }
-        for ((paramIndex, info) in allCapturedBuilder.listCaptured().filter { it.functionalArgument !is LambdaInfo }.withIndex()) {
-            if (!info.isSkipInConstructor) {
+        var nextParamOffset = 0
+        for (param in constructorParams) {
+            if (param.isSkipped) continue
+            val info = param.fieldEquivalent ?: param
+            if (info is CapturedParamInfo && !info.isSkipInConstructor) {
                 val desc = info.type.descriptor
                 val access = AsmUtil.NO_FLAG_PACKAGE_PRIVATE or Opcodes.ACC_SYNTHETIC or Opcodes.ACC_FINAL
                 classBuilder.newField(NO_ORIGIN, access, info.newFieldName, desc, null, null)
                 constructorVisitor.visitVarInsn(Opcodes.ALOAD, 0)
-                constructorVisitor.visitVarInsn(info.type.getOpcode(Opcodes.ILOAD), capturedIndexes[paramIndex])
+                constructorVisitor.visitVarInsn(info.type.getOpcode(Opcodes.ILOAD), nextParamOffset)
                 constructorVisitor.visitFieldInsn(Opcodes.PUTFIELD, transformationInfo.newClassName, info.newFieldName, desc)
             }
+            nextParamOffset += param.type.size
         }
 
-        //then transform constructor
-        //HACK: in inlining into constructor we access original captured fields with field access not local var
-        //but this fields added to general params (this assumes local var access) not captured one,
-        //so we need to add them to captured params
-        for (info in constructorAdditionalFakeParams) {
-            val fake = constructorInlineBuilder.addCapturedParamCopy(info)
-
-            if (fake.functionalArgument is LambdaInfo) {
-                //set remap value to skip this fake (captured with lambda already skipped)
-                val composed = StackValue.field(
-                    fake.type,
-                    oldObjectType,
-                    fake.newFieldName,
-                    false,
-                    StackValue.LOCAL_0
-                )
-                fake.remapValue = composed
-            }
+        for (param in constructorParams) {
+            val info = param.fieldEquivalent ?: continue
+            constructorInlineBuilder.addCapturedParamCopy(info)
         }
 
         val intermediateMethodNode = MethodNode(constructor!!.access, "<init>", constructorDescriptor, null, ArrayUtil.EMPTY_STRING_ARRAY)
@@ -457,11 +432,10 @@ class AnonymousObjectTransformer(
         constructorParamBuilder: ParametersBuilder,
         transformationInfo: AnonymousObjectTransformationInfo,
         parentFieldRemapper: FieldRemapper
-    ): List<CapturedParamInfo> {
+    ) {
+        val capturedParams = HashMap<Int, CapturedParamInfo>()
         val capturedLambdas = LinkedHashSet<LambdaInfo>() //captured var of inlined parameter
-        val constructorAdditionalFakeParams = ArrayList<CapturedParamInfo>()
         val indexToFunctionalArgument = transformationInfo.functionalArguments
-        val capturedParams = HashSet<Int>()
 
         // Possible cases where we need to add each lambda's captures separately:
         //
@@ -508,9 +482,7 @@ class AnonymousObjectTransformer(
             if (functionalArgument is LambdaInfo) {
                 capturedLambdas.add(functionalArgument)
             }
-            constructorAdditionalFakeParams.add(info)
-            capturedParams.add(varIndex)
-
+            capturedParams[varIndex] = info
             toDelete.add(parameterAload.previous)
             toDelete.add(parameterAload)
             toDelete.add(fieldNode)
@@ -521,11 +493,17 @@ class AnonymousObjectTransformer(
 
         val paramTypes = transformationInfo.constructorDesc?.let { Type.getArgumentTypes(it) } ?: emptyArray()
         for (type in paramTypes) {
-            val info = indexToFunctionalArgument[constructorParamBuilder.nextParameterOffset]
-            val isCaptured = capturedParams.contains(constructorParamBuilder.nextParameterOffset)
-            val parameterInfo = constructorParamBuilder.addNextParameter(type, info is LambdaInfo)
-            parameterInfo.functionalArgument = info
-            parameterInfo.isCaptured = isCaptured
+            val functionalArgument = indexToFunctionalArgument[constructorParamBuilder.nextParameterOffset]
+            val fieldEquivalent = capturedParams[constructorParamBuilder.nextParameterOffset]
+            val parameterInfo = constructorParamBuilder.addNextParameter(type, functionalArgument is LambdaInfo)
+            parameterInfo.functionalArgument = functionalArgument
+            parameterInfo.fieldEquivalent = fieldEquivalent
+            if (functionalArgument is LambdaInfo && parameterInfo.fieldEquivalent == null) {
+                // TODO: check if this is enough to support lambdas that have no field equivalent because they are only used
+                //  in the constructor - see `LocalClassContext` in `LocalDeclarationsLowering`.
+                // TODO: these lambdas' captures should have no fields either.
+                capturedLambdas.add(functionalArgument)
+            }
         }
 
         //For all inlined lambdas add their captured parameters
@@ -576,8 +554,6 @@ class AnonymousObjectTransformer(
 
         transformationInfo.allRecapturedParameters = allRecapturedParameters
         transformationInfo.capturedLambdasToInline = capturedLambdas.associateBy { it.lambdaClassType.internalName }
-
-        return constructorAdditionalFakeParams
     }
 
     private fun addUniqueField(name: String): String {

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

@@ -32,7 +32,14 @@ class CapturedParamInfo : ParameterInfo {
     val containingLambdaName: String
         get() = desc.containingLambdaName
 
-    constructor(desc: CapturedParamDesc, newFieldName: String, skipped: Boolean, index: Int, remapIndex: Int) : super(
+    constructor(
+        desc: CapturedParamDesc,
+        newFieldName: String,
+        skipped: Boolean,
+        index: Int,
+        remapIndex: Int,
+        skipInConstructor: Boolean
+    ) : super(
         desc.type,
         skipped,
         index,
@@ -41,10 +48,10 @@ class CapturedParamInfo : ParameterInfo {
     ) {
         this.desc = desc
         this.newFieldName = newFieldName
-        this.isSkipInConstructor = false
+        this.isSkipInConstructor = skipInConstructor
     }
 
-    constructor(
+    private constructor(
         desc: CapturedParamDesc,
         newFieldName: String,
         skipped: Boolean,

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

@@ -27,8 +27,10 @@ open class ParameterInfo(
     val declarationIndex: Int,
     val typeOnStack: Type = type
 ) {
+    // Parameters of an anonymous object constructor are all represented as locals, but some of them
+    // are also stored in fields.
+    var fieldEquivalent: CapturedParamInfo? = null
 
-    var isCaptured: Boolean = false
     var functionalArgument: FunctionalArgument? = null
 
     val isSkippedOrRemapped: Boolean

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

@@ -37,13 +37,15 @@ class ParametersBuilder private constructor() {
     }
 
     fun addCapturedParam(original: CapturedParamInfo, newFieldName: String): CapturedParamInfo {
-        val info = CapturedParamInfo(original.desc, newFieldName, original.isSkipped, nextParameterOffset, original.index)
+        val info = CapturedParamInfo(
+            original.desc, newFieldName, original.isSkipped, nextParameterOffset, original.index, original.isSkipInConstructor
+        )
         info.functionalArgument = original.functionalArgument
         return addParameter(info)
     }
 
     fun addCapturedParam(desc: CapturedParamDesc, newFieldName: String, skipInConstructor: Boolean): CapturedParamInfo {
-        return addParameter(CapturedParamInfo(desc, newFieldName, false, nextParameterOffset, null, skipInConstructor, -1))
+        return addParameter(CapturedParamInfo(desc, newFieldName, false, nextParameterOffset, -1, skipInConstructor))
     }
 
     fun addCapturedParamCopy(copyFrom: CapturedParamInfo): CapturedParamInfo {
@@ -59,7 +61,7 @@ class ParametersBuilder private constructor() {
         original: ParameterInfo?
     ): CapturedParamInfo {
         val info = CapturedParamInfo(
-            CapturedParamDesc(containingLambdaType, fieldName, type), newFieldName, skipped, nextParameterOffset, original?.index ?: -1
+            CapturedParamDesc(containingLambdaType, fieldName, type), newFieldName, skipped, nextParameterOffset, original?.index ?: -1, false
         )
         if (original != null) {
             info.functionalArgument = original.functionalArgument