KT-52027 Encapsulate and fix calculation of parameter indices

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

@@ -230,7 +230,10 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
             return processDefaultMaskOrMethodHandler(stackValue, kind)
         }
 
-        val info = invocationParamBuilder.addNextValueParameter(jvmKotlinType.type, false, null, parameterIndex)
+        val info = when (parameterIndex) {
+            -1 -> invocationParamBuilder.addNextParameter(jvmKotlinType.type, false)
+            else -> invocationParamBuilder.addNextValueParameter(jvmKotlinType.type, false, null, parameterIndex)
+        }
         info.functionalArgument = when (kind) {
             ValueKind.READ_OF_INLINE_LAMBDA_FOR_INLINE_SUSPEND_PARAMETER ->
                 NonInlineArgumentForInlineSuspendParameter.INLINE_LAMBDA_AS_VARIABLE

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

@@ -20,19 +20,17 @@ class ParametersBuilder private constructor() {
     private var nextValueParameterIndex = 0
 
     fun addThis(type: Type, skipped: Boolean): ParameterInfo {
-        return addParameter(ParameterInfo(type, skipped, nextParameterOffset, -1, nextValueParameterIndex))
+        return addParameter(ParameterInfo(type, skipped, nextParameterOffset, -1, nextValueParameterIndex), false)
     }
 
     fun addNextParameter(type: Type, skipped: Boolean, typeOnStack: Type = type): ParameterInfo {
-        return addParameter(ParameterInfo(type, skipped, nextParameterOffset, null, nextValueParameterIndex, typeOnStack))
+        return addParameter(ParameterInfo(type, skipped, nextParameterOffset, null, nextValueParameterIndex, typeOnStack), false)
     }
 
     fun addNextValueParameter(type: Type, skipped: Boolean, remapValue: StackValue?, parameterIndex: Int): ParameterInfo {
         return addParameter(
-            ParameterInfo(
+            ParameterInfo(type, skipped, nextParameterOffset, remapValue, parameterIndex + valueParamFirstIndex),
-                type, skipped, nextParameterOffset, remapValue,
+            true
-                if (parameterIndex == -1) nextValueParameterIndex else parameterIndex + valueParamFirstIndex
-            )
         )
     }
 
@@ -41,15 +39,15 @@ class ParametersBuilder private constructor() {
             original.desc, newFieldName, original.isSkipped, nextParameterOffset, original.index, original.isSkipInConstructor
         )
         info.functionalArgument = original.functionalArgument
-        return addParameter(info)
+        return addParameter(info, false)
     }
 
     fun addCapturedParam(desc: CapturedParamDesc, newFieldName: String, skipInConstructor: Boolean): CapturedParamInfo {
-        return addParameter(CapturedParamInfo(desc, newFieldName, false, nextParameterOffset, -1, skipInConstructor))
+        return addParameter(CapturedParamInfo(desc, newFieldName, false, nextParameterOffset, -1, skipInConstructor), false)
     }
 
     fun addCapturedParamCopy(copyFrom: CapturedParamInfo): CapturedParamInfo {
-        return addParameter(copyFrom.cloneWithNewDeclarationIndex(-1))
+        return addParameter(copyFrom.cloneWithNewDeclarationIndex(-1), false)
     }
 
     fun addCapturedParam(
@@ -66,23 +64,21 @@ class ParametersBuilder private constructor() {
         if (original != null) {
             info.functionalArgument = original.functionalArgument
         }
-        return addParameter(info)
+        return addParameter(info, false)
     }
 
-    private fun <T : ParameterInfo> addParameter(info: T): T {
+    private fun <T : ParameterInfo> addParameter(info: T, isValueParameter: Boolean): T {
         params.add(info)
         nextParameterOffset += info.type.size
         if (info !is CapturedParamInfo) {
             nextValueParameterIndex++
+            if (!isValueParameter) {
+                valueParamFirstIndex++
+            }
         }
         return info
     }
 
-    fun markValueParametersStart(contextReceiversCount: Int) {
-        this.valueParamFirstIndex = params.size - contextReceiversCount
-        this.nextValueParameterIndex = valueParamFirstIndex
-    }
-
     fun listCaptured(): List<CapturedParamInfo> {
         return params.filterIsInstance<CapturedParamInfo>()
     }

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

@@ -97,8 +97,7 @@ class PsiInlineCodegen(
             hiddenParameters += invocationParamBuilder.addNextParameter(param.asmType, false) to
                     codegen.frameMap.enterTemp(param.asmType)
         }
-        // TODO: Add context receivers as hiddenParameters and pass their count
+        // TODO: Add context receivers as hiddenParameters
-        invocationParamBuilder.markValueParametersStart(0)
     }
 
     override fun putHiddenParamsIntoLocals() {

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -17462,6 +17462,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51863.kt");
             }
 
+            @Test
+            @TestMetadata("kt52207.kt")
+            public void testKt52207() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt52207.kt");
+            }
+
             @Test
             @TestMetadata("overload.kt")
             public void testOverload() throws Exception {

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

@@ -509,7 +509,6 @@ class ExpressionCodegen(
             callGenerator.genValueAndPut(callee.extensionReceiverParameter!!, receiver, type, this, data)
         }
 
-        callGenerator.beforeValueParametersStart(callee.contextReceiverParametersCount)
         callee.valueParameters.subList(callee.contextReceiverParametersCount, callee.valueParameters.size)
             .forEachIndexed { i, valueParameter -> handleValueParameter(i + contextReceivers.size, valueParameter) }
 

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

@@ -36,8 +36,6 @@ interface IrCallGenerator {
 
     fun beforeCallStart() {}
 
-    fun beforeValueParametersStart(contextReceiversCount: Int) {}
-
     fun afterCallEnd() {}
 
     fun genValueAndPut(

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

@@ -150,10 +150,6 @@ class IrInlineCodegen(
         }
     }
 
-    override fun beforeValueParametersStart(contextReceiversCount: Int) {
-        invocationParamBuilder.markValueParametersStart(contextReceiversCount)
-    }
-
     override fun genInlineCall(
         callableMethod: IrCallableMethod,
         codegen: ExpressionCodegen,

compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt52207.kt

@@ -0,0 +1,20 @@
+// !LANGUAGE: +ContextReceivers
+// TARGET_BACKEND: JVM_IR
+
+class A
+
+class Example {
+    context(A)
+    inline fun fn(x: Int) {}
+}
+
+fun test() {
+    with(A()) {
+        Example().fn(1)
+    }
+}
+
+fun box(): String {
+    test()
+    return "OK"
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -17462,6 +17462,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51863.kt");
             }
 
+            @Test
+            @TestMetadata("kt52207.kt")
+            public void testKt52207() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt52207.kt");
+            }
+
             @Test
             @TestMetadata("overload.kt")
             public void testOverload() throws Exception {