KT-51243 Fix parameterized contextual lambda

^KT-51243 Fixed

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -10907,6 +10907,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/overloading.kt");
                 }
 
+                @Test
+                @TestMetadata("parameterizedContextualLambda.kt")
+                public void testParameterizedContextualLambda() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/parameterizedContextualLambda.kt");
+                }
+
                 @Test
                 @TestMetadata("plusMatrix.kt")
                 public void testPlusMatrix() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -10907,6 +10907,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/overloading.kt");
                 }
 
+                @Test
+                @TestMetadata("parameterizedContextualLambda.kt")
+                public void testParameterizedContextualLambda() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/parameterizedContextualLambda.kt");
+                }
+
                 @Test
                 @TestMetadata("plusMatrix.kt")
                 public void testPlusMatrix() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java

@@ -10907,6 +10907,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/overloading.kt");
                 }
 
+                @Test
+                @TestMetadata("parameterizedContextualLambda.kt")
+                public void testParameterizedContextualLambda() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/parameterizedContextualLambda.kt");
+                }
+
                 @Test
                 @TestMetadata("plusMatrix.kt")
                 public void testPlusMatrix() throws Exception {

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

@@ -16734,6 +16734,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/overloading.kt");
             }
 
+            @Test
+            @TestMetadata("parameterizedContextualLambda.kt")
+            public void testParameterizedContextualLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/parameterizedContextualLambda.kt");
+            }
+
             @Test
             @TestMetadata("passingLambdaToContextualParam.kt")
             public void testPassingLambdaToContextualParam() throws Exception {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedAtoms.kt

@@ -19,11 +19,11 @@ import org.jetbrains.kotlin.resolve.calls.model.LambdaWithTypeVariableAsExpected
 import org.jetbrains.kotlin.resolve.calls.model.PostponedCallableReferenceMarker
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.utils.addIfNotNull
 
 //  --------------------------- Variables ---------------------------
 
 
-
 //  -------------------------- Atoms --------------------------
 
 sealed class PostponedResolvedAtom : PostponedResolvedAtomMarker {
@@ -62,7 +62,16 @@ class ResolvedLambdaAtom(
 
     lateinit var returnStatements: Collection<FirStatement>
 
-    override val inputTypes: Collection<ConeKotlinType> get() = receiver?.let { parameters + it } ?: parameters
+    override val inputTypes: Collection<ConeKotlinType>
+        get() {
+            if (receiver == null && contextReceivers.isEmpty()) return parameters
+            return ArrayList<ConeKotlinType>(parameters.size + contextReceivers.size + (if (receiver != null) 1 else 0)).apply {
+                addAll(parameters)
+                addIfNotNull(receiver)
+                addAll(contextReceivers)
+            }
+        }
+
     override val outputType: ConeKotlinType get() = returnType
 
     fun replaceExpectedType(expectedType: ConeKotlinType, newReturnType: ConeTypeVariableType) {

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolutionAtoms.kt

@@ -8,20 +8,26 @@ package org.jetbrains.kotlin.resolve.calls.model
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.resolve.calls.components.*
+import org.jetbrains.kotlin.resolve.calls.components.ReturnArgumentsInfo
+import org.jetbrains.kotlin.resolve.calls.components.TypeArgumentsToParametersMapper
+import org.jetbrains.kotlin.resolve.calls.components.candidate.CallableReferenceResolutionCandidate
+import org.jetbrains.kotlin.resolve.calls.components.candidate.ResolutionCandidate
+import org.jetbrains.kotlin.resolve.calls.components.extractInputOutputTypesFromCallableReferenceExpectedType
+import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.components.FreshVariableNewTypeSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
-import org.jetbrains.kotlin.resolve.calls.inference.model.*
+import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintError
+import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintMismatch
+import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintWarning
+import org.jetbrains.kotlin.resolve.calls.inference.model.TypeVariableForLambdaReturnType
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
-import org.jetbrains.kotlin.resolve.calls.components.candidate.ResolutionCandidate
-import org.jetbrains.kotlin.resolve.calls.components.candidate.CallableReferenceResolutionCandidate
-import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.constants.IntegerValueTypeConstant
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeConstructor
 import org.jetbrains.kotlin.types.UnwrappedType
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.typeUtil.unCapture
+import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 /**
@@ -174,7 +180,16 @@ class ResolvedLambdaAtom(
         setAnalyzedResults(subResolvedAtoms)
     }
 
-    override val inputTypes: Collection<UnwrappedType> get() = receiver?.let { parameters + it } ?: parameters
+    override val inputTypes: Collection<UnwrappedType>
+        get() {
+            if (receiver == null && contextReceivers.isEmpty()) return parameters
+            return ArrayList<UnwrappedType>(parameters.size + contextReceivers.size + (if (receiver != null) 1 else 0)).apply {
+                addAll(parameters)
+                addIfNotNull(receiver)
+                addAll(contextReceivers)
+            }
+        }
+
     override val outputType: UnwrappedType get() = returnType
 }
 
@@ -230,8 +245,7 @@ class CallableReferenceWithRevisedExpectedTypeAtom(
 class PostponedCallableReferenceAtom(
     eagerCallableReferenceAtom: EagerCallableReferenceAtom
 ) : AbstractPostponedCallableReferenceAtom(eagerCallableReferenceAtom.atom, eagerCallableReferenceAtom.expectedType),
-    PostponedCallableReferenceMarker
-{
+    PostponedCallableReferenceMarker {
     override var revisedExpectedType: UnwrappedType? = null
         private set
 

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

@@ -0,0 +1,13 @@
+// !LANGUAGE: +ContextReceivers
+// TARGET_BACKEND: JVM_IR
+// FIR status: context receivers aren't yet supported
+
+class A {
+    val result = "OK"
+}
+
+fun <T> test(receiver: T, action: context(T) () -> String) = action(receiver)
+
+fun box(): String = with(A()) {
+    result
+}

compiler/testData/diagnostics/tests/extensions/contextReceivers/parameterizedContextualLambda.kt

@@ -0,0 +1,20 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+// !LANGUAGE: +ContextReceivers
+
+fun <T> test(action: context(T) () -> Unit) {}
+
+fun <T> test2(actionWithArg: context(T) (T) -> Unit) {}
+
+fun main() {
+    test<String> {
+        length
+    }
+    test<Int> {
+        toDouble()
+    }
+    test2<String> { a ->
+        length
+        a.length
+    }
+}

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

@@ -10913,6 +10913,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/overloading.kt");
                 }
 
+                @Test
+                @TestMetadata("parameterizedContextualLambda.kt")
+                public void testParameterizedContextualLambda() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/extensions/contextReceivers/parameterizedContextualLambda.kt");
+                }
+
                 @Test
                 @TestMetadata("plusMatrix.kt")
                 public void testPlusMatrix() throws Exception {

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

@@ -16734,6 +16734,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/overloading.kt");
             }
 
+            @Test
+            @TestMetadata("parameterizedContextualLambda.kt")
+            public void testParameterizedContextualLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/parameterizedContextualLambda.kt");
+            }
+
             @Test
             @TestMetadata("passingLambdaToContextualParam.kt")
             public void testPassingLambdaToContextualParam() throws Exception {