JVM: pregenerate inline lambda bodies as early as possible

Which in JVM_IR is immediately in `genValueAndPut`, but for the old
backend needs to be delayed until `genCallInner` for some reason.

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

@@ -48,7 +48,7 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
     protected val jvmSignature: JvmMethodSignature,
     private val typeParameterMappings: TypeParameterMappings<*>,
     protected val sourceCompiler: SourceCompilerForInline,
-    private val reifiedTypeInliner: ReifiedTypeInliner<*>
+    protected val reifiedTypeInliner: ReifiedTypeInliner<*>
 ) {
     init {
         assert(InlineUtil.isInline(functionDescriptor)) {
@@ -193,21 +193,27 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
                 ?: error("No stack value for continuation parameter of suspend function")
     }
 
+    protected fun rememberClosure(parameterType: Type, index: Int, lambdaInfo: LambdaInfo) {
+        val closureInfo = invocationParamBuilder.addNextValueParameter(parameterType, true, null, index)
+        closureInfo.functionalArgument = lambdaInfo
+        expressionMap[closureInfo.index] = lambdaInfo
+    }
+
     private fun inlineCall(nodeAndSmap: SMAPAndMethodNode, inlineDefaultLambda: Boolean): InlineResult {
         assert(delayedHiddenWriting == null) { "'putHiddenParamsIntoLocals' should be called after 'processAndPutHiddenParameters(true)'" }
+
         val node = nodeAndSmap.node
         if (inlineDefaultLambda) {
             for (lambda in extractDefaultLambdas(node)) {
                 invocationParamBuilder.buildParameters().getParameterByDeclarationSlot(lambda.offset).functionalArgument = lambda
                 val prev = expressionMap.put(lambda.offset, lambda)
                 assert(prev == null) { "Lambda with offset ${lambda.offset} already exists: $prev" }
+                lambda.generateLambdaBody(sourceCompiler, reifiedTypeInliner)
+                rememberCapturedForDefaultLambda(lambda)
             }
         }
-        val reificationResult = reifiedTypeInliner.reifyInstructions(node)
-        generateClosuresBodies()
 
-        //through generation captured parameters will be added to invocationParamBuilder
+        val reificationResult = reifiedTypeInliner.reifyInstructions(node)
-        putClosureParametersOnStack()
 
         val parameters = invocationParamBuilder.buildParameters()
 
@@ -340,14 +346,6 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
         }
     }
 
-    private fun generateClosuresBodies() {
-        for (info in expressionMap.values) {
-            if (info is LambdaInfo) {
-                info.generateLambdaBody(sourceCompiler, reifiedTypeInliner)
-            }
-        }
-    }
-
     protected fun putArgumentOrCapturedToLocalVal(
         jvmKotlinType: JvmKotlinType,
         stackValue: StackValue,
@@ -431,19 +429,7 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
         }
     }
 
-    private fun putClosureParametersOnStack() {
+    private fun rememberCapturedForDefaultLambda(defaultLambda: DefaultLambda) {
-        for (next in expressionMap.values) {
-            if (next is LambdaInfo) {
-                // Closure parameters for bound references have already been generated in between other arguments.
-                if (next is ExpressionLambda && next.isBoundCallableReference) continue
-                putClosureParametersOnStack(next)
-            }
-        }
-    }
-
-    protected abstract fun putClosureParametersOnStack(next: LambdaInfo)
-
-    protected fun rememberCapturedForDefaultLambda(defaultLambda: DefaultLambda) {
         for ((paramIndex, captured) in defaultLambda.capturedVars.withIndex()) {
             putArgumentOrCapturedToLocalVal(
                 JvmKotlinType(captured.type),
@@ -459,7 +445,6 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
         }
     }
 
-
     protected fun processDefaultMaskOrMethodHandler(value: StackValue, kind: ValueKind): Boolean {
         if (kind !== ValueKind.DEFAULT_MASK && kind !== ValueKind.METHOD_HANDLE_IN_DEFAULT) {
             return false

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

@@ -28,7 +28,6 @@ import org.jetbrains.kotlin.resolve.inline.InlineUtil
 import org.jetbrains.kotlin.resolve.inline.InlineUtil.isInlinableParameterExpression
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterKind
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
-import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DescriptorWithContainerSource
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -72,6 +71,17 @@ class PsiInlineCodegen(
         }
         try {
             val registerLineNumber = registerLineNumberAfterwards(resolvedCall)
+            for (info in expressionMap.values) {
+                if (info is PsiExpressionLambda) {
+                    // Can't be done immediately in `rememberClosure` for some reason:
+                    info.generateLambdaBody(sourceCompiler, reifiedTypeInliner)
+                    // Requires `generateLambdaBody` first if the closure is non-empty (for bound callable references,
+                    // or indeed any callable references, it *is* empty, so this was done in `rememberClosure`):
+                    if (!info.isBoundCallableReference) {
+                        putClosureParametersOnStack(info, null)
+                    }
+                }
+            }
             performInline(resolvedCall?.typeArguments?.keys?.toList(), callDefault, callDefault, codegen.typeSystem, registerLineNumber)
         } finally {
             state.globalInlineContext.exitFromInlining()
@@ -102,31 +112,6 @@ class PsiInlineCodegen(
         delayedHiddenWriting = recordParameterValueInLocalVal(justProcess, false, *hiddenParameters.toTypedArray())
     }
 
-    var activeLambda: LambdaInfo? = null
-        private set
-
-    override fun putClosureParametersOnStack(next: LambdaInfo) {
-        activeLambda = next
-        when (next) {
-            is PsiExpressionLambda -> {
-                val receiverValue = next.boundReceiver?.let { boundReceiver ->
-                    val receiver = codegen.generateReceiverValue(boundReceiver, false)
-                    val receiverKotlinType = receiver.kotlinType
-                    val boxedReceiver =
-                        if (receiverKotlinType != null)
-                            DescriptorAsmUtil.boxType(receiver.type, receiverKotlinType, state.typeMapper)
-                        else
-                            AsmUtil.boxType(receiver.type)
-                    StackValue.coercion(receiver, boxedReceiver, receiverKotlinType)
-                }
-                codegen.pushClosureOnStack(next.classDescriptor, true, this, receiverValue)
-            }
-            is DefaultLambda -> rememberCapturedForDefaultLambda(next)
-            else -> throw RuntimeException("Unknown lambda: $next")
-        }
-        activeLambda = null
-    }
-
     /*lambda or callable reference*/
     private fun isInliningParameter(expression: KtExpression, valueParameterDescriptor: ValueParameterDescriptor): Boolean {
         //TODO deparenthesize typed
@@ -147,11 +132,7 @@ class PsiInlineCodegen(
         }
 
         if (isInliningParameter(argumentExpression, valueParameterDescriptor)) {
-            val lambdaInfo = rememberClosure(argumentExpression, parameterType.type, valueParameterDescriptor)
+            rememberClosure(argumentExpression, parameterType.type, valueParameterDescriptor)
-            if (lambdaInfo.boundReceiver != null) {
-                // Has to be done immediately to preserve evaluation order.
-                putClosureParametersOnStack(lambdaInfo)
-            }
         } else {
             val value = codegen.gen(argumentExpression)
             val kind = when {
@@ -170,7 +151,7 @@ class PsiInlineCodegen(
     private fun isCallSiteIsSuspend(descriptor: ValueParameterDescriptor): Boolean =
         state.bindingContext[CodegenBinding.CALL_SITE_IS_SUSPEND_FOR_CROSSINLINE_LAMBDA, descriptor] == true
 
-    private fun rememberClosure(expression: KtExpression, type: Type, parameter: ValueParameterDescriptor): PsiExpressionLambda {
+    private fun rememberClosure(expression: KtExpression, type: Type, parameter: ValueParameterDescriptor) {
         val ktLambda = KtPsiUtil.deparenthesize(expression)
         assert(isInlinableParameterExpression(ktLambda)) { "Couldn't find inline expression in ${expression.text}" }
 
@@ -179,15 +160,33 @@ class PsiInlineCodegen(
             JvmCodegenUtil.getBoundCallableReferenceReceiver(resolvedCall)
         } else null
 
-        return PsiExpressionLambda(
+        val lambda = PsiExpressionLambda(
-            ktLambda!!, state.typeMapper, state.languageVersionSettings, parameter.isCrossinline, boundReceiver
+            ktLambda!!, state.typeMapper, state.languageVersionSettings, parameter.isCrossinline, boundReceiver != null
-        ).also { lambda ->
+        )
-            val closureInfo = invocationParamBuilder.addNextValueParameter(type, true, null, parameter.index)
+        rememberClosure(type, parameter.index, lambda)
-            closureInfo.functionalArgument = lambda
+        if (boundReceiver != null) {
-            expressionMap[closureInfo.index] = lambda
+            // Has to be done immediately to preserve evaluation order.
+            val receiver = codegen.generateReceiverValue(boundReceiver, false)
+            val receiverKotlinType = receiver.kotlinType
+            val boxedReceiver =
+                if (receiverKotlinType != null)
+                    DescriptorAsmUtil.boxType(receiver.type, receiverKotlinType, state.typeMapper)
+                else
+                    AsmUtil.boxType(receiver.type)
+            val receiverValue = StackValue.coercion(receiver, boxedReceiver, receiverKotlinType)
+            putClosureParametersOnStack(lambda, receiverValue)
         }
     }
 
+    var activeLambda: LambdaInfo? = null
+        private set
+
+    private fun putClosureParametersOnStack(next: PsiExpressionLambda, receiverValue: StackValue?) {
+        activeLambda = next
+        codegen.pushClosureOnStack(next.classDescriptor, true, this, receiverValue)
+        activeLambda = null
+    }
+
     override fun putValueIfNeeded(parameterType: JvmKotlinType, value: StackValue, kind: ValueKind, parameterIndex: Int) {
         if (processDefaultMaskOrMethodHandler(value, kind)) return
 
@@ -231,11 +230,8 @@ class PsiExpressionLambda(
     private val typeMapper: KotlinTypeMapper,
     private val languageVersionSettings: LanguageVersionSettings,
     isCrossInline: Boolean,
-    val boundReceiver: ReceiverValue?
-) : ExpressionLambda(isCrossInline) {
     override val isBoundCallableReference: Boolean
-        get() = boundReceiver != null
+) : ExpressionLambda(isCrossInline) {
-
     override val lambdaClassType: Type
 
     override val invokeMethod: Method
@@ -297,6 +293,7 @@ class PsiExpressionLambda(
         isSuspend = invokeMethodDescriptor.isSuspend
     }
 
+    // This can only be computed after generating the body, hence `lazy`.
     override val capturedVars: List<CapturedParamDesc> by lazy {
         arrayListOf<CapturedParamDesc>().apply {
             val captureThis = closure.capturedOuterClassDescriptor

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

@@ -57,16 +57,6 @@ class IrInlineCodegen(
         }
     }
 
-    override fun putClosureParametersOnStack(next: LambdaInfo) {
-        when (next) {
-            is IrExpressionLambdaImpl -> next.reference.getArgumentsWithIr().forEachIndexed { index, (_, ir) ->
-                putCapturedValueOnStack(ir, next.capturedVars[index], index)
-            }
-            is IrDefaultLambda -> rememberCapturedForDefaultLambda(next)
-            else -> throw RuntimeException("Unknown lambda: $next")
-        }
-    }
-
     override fun genValueAndPut(
         irValueParameter: IrValueParameter,
         argumentExpression: IrExpression,
@@ -78,12 +68,10 @@ class IrInlineCodegen(
         if (isInlineParameter && argumentExpression.isInlineIrExpression()) {
             val irReference = (argumentExpression as IrBlock).statements.filterIsInstance<IrFunctionReference>().single()
             val lambdaInfo = IrExpressionLambdaImpl(codegen, irReference, irValueParameter)
-            val closureInfo = invocationParamBuilder.addNextValueParameter(parameterType, true, null, irValueParameter.index)
+            rememberClosure(parameterType, irValueParameter.index, lambdaInfo)
-            closureInfo.functionalArgument = lambdaInfo
+            lambdaInfo.generateLambdaBody(sourceCompiler, reifiedTypeInliner)
-            expressionMap[closureInfo.index] = lambdaInfo
+            lambdaInfo.reference.getArgumentsWithIr().forEachIndexed { index, (_, ir) ->
-            val boundReceiver = irReference.extensionReceiver
+                putCapturedValueOnStack(ir, lambdaInfo.capturedVars[index], index)
-            if (boundReceiver != null) {
-                putCapturedValueOnStack(boundReceiver, lambdaInfo.capturedVars.single(), 0)
             }
         } else {
             val kind = when (irValueParameter.origin) {