JVM: move activeLambda to PsiInlineCodegen

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -1261,8 +1261,8 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             if (closure.isSuspendLambda()) {
                 // When inlining crossinline lambda, the ACONST_NULL is never popped.
                 // Thus, do not generate it. Otherwise, it leads to VerifyError on run-time.
-                boolean isCrossinlineLambda = (callGenerator instanceof InlineCodegen<?>) &&
+                boolean isCrossinlineLambda = (callGenerator instanceof PsiInlineCodegen) &&
-                                              Objects.requireNonNull(((InlineCodegen) callGenerator).getActiveLambda(),
+                                              Objects.requireNonNull(((PsiInlineCodegen) callGenerator).getActiveLambda(),
                                                                      "no active lambda found").isCrossInline;
                 if (!isCrossinlineLambda) {
                     v.aconst(null);

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

@@ -69,9 +69,6 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
 
     protected val expressionMap = linkedMapOf<Int, FunctionalArgument>()
 
-    var activeLambda: LambdaInfo? = null
-        protected set
-
     private val sourceMapper = sourceCompiler.lazySourceMapper
 
     protected var delayedHiddenWriting: Function0<Unit>? = null
@@ -356,7 +353,7 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
     protected fun putArgumentOrCapturedToLocalVal(
         jvmKotlinType: JvmKotlinType,
         stackValue: StackValue,
-        capturedParamIndex: Int,
+        capturedParam: CapturedParamDesc?,
         parameterIndex: Int,
         kind: ValueKind
     ) {
@@ -373,9 +370,8 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
             val remappedValue = if (couldBeRemapped) stackValue else null
 
             val info: ParameterInfo
-            if (capturedParamIndex >= 0) {
+            if (capturedParam != null) {
-                val capturedParamInfoInLambda = activeLambda!!.capturedVars[capturedParamIndex]
+                info = invocationParamBuilder.addCapturedParam(capturedParam, capturedParam.fieldName, false)
-                info = invocationParamBuilder.addCapturedParam(capturedParamInfoInLambda, capturedParamInfoInLambda.fieldName, false)
                 info.remapValue = remappedValue
             } else {
                 info = invocationParamBuilder.addNextValueParameter(jvmType, false, remappedValue, parameterIndex)
@@ -456,7 +452,7 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
                 //HACK: actually parameter would be placed on stack in default function
                 // also see ValueKind.DEFAULT_LAMBDA_CAPTURED_PARAMETER check
                 StackValue.onStack(captured.type),
-                paramIndex,
+                captured,
                 paramIndex,
                 ValueKind.DEFAULT_LAMBDA_CAPTURED_PARAMETER
             )

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

@@ -102,6 +102,9 @@ class PsiInlineCodegen(
         delayedHiddenWriting = recordParameterValueInLocalVal(justProcess, false, *hiddenParameters.toTypedArray())
     }
 
+    var activeLambda: LambdaInfo? = null
+        private set
+
     override fun putClosureParametersOnStack(next: LambdaInfo, functionReferenceReceiver: StackValue?) {
         activeLambda = next
         when (next) {
@@ -192,12 +195,13 @@ class PsiInlineCodegen(
 
         assert(maskValues.isEmpty()) { "Additional default call arguments should be last ones, but $value" }
 
-        putArgumentOrCapturedToLocalVal(parameterType, value, -1, parameterIndex, kind)
+        putArgumentOrCapturedToLocalVal(parameterType, value, null, parameterIndex, kind)
     }
 
     override fun putCapturedValueOnStack(stackValue: StackValue, valueType: Type, paramIndex: Int) {
+        val param = activeLambda!!.capturedVars[paramIndex]
         putArgumentOrCapturedToLocalVal(
-            JvmKotlinType(stackValue.type, stackValue.kotlinType), stackValue, paramIndex, paramIndex, ValueKind.CAPTURED
+            JvmKotlinType(stackValue.type, stackValue.kotlinType), stackValue, param, paramIndex, ValueKind.CAPTURED
         )
     }
 

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

@@ -58,17 +58,13 @@ class IrInlineCodegen(
     }
 
     override fun putClosureParametersOnStack(next: LambdaInfo, functionReferenceReceiver: StackValue?) {
-        activeLambda = next
-
         when (next) {
             is IrExpressionLambdaImpl -> next.reference.getArgumentsWithIr().forEachIndexed { index, (_, ir) ->
-                putCapturedValueOnStack(ir, next.capturedVars[index].type, index)
+                putCapturedValueOnStack(ir, next.capturedVars[index], index)
             }
             is IrDefaultLambda -> rememberCapturedForDefaultLambda(next)
             else -> throw RuntimeException("Unknown lambda: $next")
         }
-
-        activeLambda = null
     }
 
     override fun genValueAndPut(
@@ -87,9 +83,7 @@ class IrInlineCodegen(
             expressionMap[closureInfo.index] = lambdaInfo
             val boundReceiver = irReference.extensionReceiver
             if (boundReceiver != null) {
-                activeLambda = lambdaInfo
+                putCapturedValueOnStack(boundReceiver, lambdaInfo.capturedVars.single(), 0)
-                putCapturedValueOnStack(boundReceiver, lambdaInfo.capturedVars.single().type, 0)
-                activeLambda = null
             }
         } else {
             val kind = when (irValueParameter.origin) {
@@ -125,15 +119,15 @@ class IrInlineCodegen(
             //TODO support default argument erasure
             if (!processDefaultMaskOrMethodHandler(onStack, kind)) {
                 val expectedType = JvmKotlinType(parameterType, irValueParameter.type.toIrBasedKotlinType())
-                putArgumentOrCapturedToLocalVal(expectedType, onStack, -1, irValueParameter.index, kind)
+                putArgumentOrCapturedToLocalVal(expectedType, onStack, null, irValueParameter.index, kind)
             }
         }
     }
 
-    private fun putCapturedValueOnStack(argumentExpression: IrExpression, valueType: Type, capturedParamIndex: Int) {
+    private fun putCapturedValueOnStack(argumentExpression: IrExpression, param: CapturedParamDesc, capturedParamIndex: Int) {
-        val onStack = codegen.genOrGetLocal(argumentExpression, valueType, argumentExpression.type, BlockInfo())
+        val onStack = codegen.genOrGetLocal(argumentExpression, param.type, argumentExpression.type, BlockInfo())
-        val expectedType = JvmKotlinType(valueType, argumentExpression.type.toIrBasedKotlinType())
+        val expectedType = JvmKotlinType(param.type, argumentExpression.type.toIrBasedKotlinType())
-        putArgumentOrCapturedToLocalVal(expectedType, onStack, capturedParamIndex, capturedParamIndex, ValueKind.CAPTURED)
+        putArgumentOrCapturedToLocalVal(expectedType, onStack, param, capturedParamIndex, ValueKind.CAPTURED)
     }
 
     override fun beforeValueParametersStart() {