[FIR] Fix extracting parameters from extension function types

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArguments.kt

@@ -17,11 +17,18 @@ import org.jetbrains.kotlin.fir.symbols.invoke
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassTypeImpl
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
-import org.jetbrains.kotlin.resolve.calls.inference.model.ArgumentConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
+fun FirTypeRef.isExtensionFunctionType(): Boolean {
+    return annotations.any {
+        (it.annotationTypeRef as? FirResolvedTypeRef)?.let {
+            it.type.toString() == "kotlin/ExtensionFunctionType"
+        } == true
+    }
+}
+
 fun Candidate.preprocessLambdaArgument(
     csBuilder: ConstraintSystemBuilder,
     argument: FirAnonymousFunction,
@@ -152,11 +159,11 @@ internal fun extractLambdaInfoFromFunctionalType(
     if (expectedType == null) return null
     if (expectedType is ConeFlexibleType) return extractLambdaInfoFromFunctionalType(expectedType.lowerBound, expectedTypeRef, argument)
     if (!expectedType.isBuiltinFunctionalType) return null
-    val parameters = extractLambdaParameters(expectedType, argument)
 
     val argumentAsFunctionExpression = argument//.safeAs<FunctionExpression>()
     val receiverType = argumentAsFunctionExpression.receiverType ?: expectedType.receiverType(expectedTypeRef)
     val returnType = argumentAsFunctionExpression.returnType ?: expectedType.returnType ?: return null
+    val parameters = extractLambdaParameters(expectedType, argument, expectedTypeRef.isExtensionFunctionType())
 
     return ResolvedLambdaAtom(
         argument,
@@ -168,9 +175,9 @@ internal fun extractLambdaInfoFromFunctionalType(
     )
 }
 
-private fun extractLambdaParameters(expectedType: ConeKotlinType, argument: FirAnonymousFunction): List<ConeKotlinType> {
+private fun extractLambdaParameters(expectedType: ConeKotlinType, argument: FirAnonymousFunction, expectedTypeIsExtensionFunctionType: Boolean): List<ConeKotlinType> {
     val parameters = argument.valueParameters
-    val expectedParameters = expectedType.valueParameterTypes
+    val expectedParameters = expectedType.extractParametersForFunctionalType(expectedTypeIsExtensionFunctionType)
 
     val nullableAnyType = StandardClassIds.Any(argument.session.firSymbolProvider).constructType(emptyArray(), true)
 
@@ -184,8 +191,17 @@ private fun extractLambdaParameters(expectedType: ConeKotlinType, argument: FirA
     }
 }
 
-class TypeVariableForLambdaReturnType(val argument: FirAnonymousFunction, name: String) : ConeTypeVariable(name)
+private fun ConeKotlinType.extractParametersForFunctionalType(isExtensionFunctionType: Boolean): List<ConeKotlinType?> {
+    return valueParameterTypes.let {
+        if (isExtensionFunctionType) {
+            it.drop(1)
+        } else {
+            it
+        }
+    }
+}
 
+class TypeVariableForLambdaReturnType(val argument: FirAnonymousFunction, name: String) : ConeTypeVariable(name)
 
 class ResolvedLambdaAtom(
     val atom: FirAnonymousFunction,

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

@@ -108,12 +108,12 @@ class FirCallCompleter(
             stubsForPostponedVariables: Map<TypeVariableMarker, StubTypeMarker>
         ): Pair<List<FirExpression>, InferenceSession> {
 
-            val needItParam = lambdaArgument.valueParameters.isEmpty() && parameters.size == (if (receiverType != null) 2 else 1)
+            val needItParam = lambdaArgument.valueParameters.isEmpty() && parameters.size == 1
 
             val itParam = when {
                 needItParam -> {
                     val name = Name.identifier("it")
-                    val itType = if (receiverType != null) parameters[1] else parameters.single()
+                    val itType = parameters.single()
                     FirValueParameterImpl(
                         null,
                         session,

compiler/testData/ir/irText/classes/lambdaInDataClassDefaultParameter.fir.txt

@@ -4,8 +4,9 @@ FILE fqName:<root> fileName:/lambdaInDataClassDefaultParameter.kt
     CONSTRUCTOR visibility:public <> (runA:kotlin.Function2<<root>.A, kotlin.String, kotlin.Unit>) returnType:<root>.A [primary]
       VALUE_PARAMETER name:runA index:0 type:kotlin.Function2<<root>.A, kotlin.String, kotlin.Unit>
         EXPRESSION_BODY
-          FUN_EXPR type=kotlin.Function1<<root>.A, kotlin.Unit> origin=LAMBDA
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.Unit
+          FUN_EXPR type=kotlin.Function2<<root>.A, kotlin.String, kotlin.Unit> origin=LAMBDA
+            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.String) returnType:kotlin.Unit
+              VALUE_PARAMETER name:it index:0 type:kotlin.String
               BLOCK_BODY
                 GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
       BLOCK_BODY

compiler/testData/ir/irText/declarations/parameters/lambdas.fir.txt

@@ -15,10 +15,12 @@ FILE fqName:<root> fileName:/lambdas.kt
   PROPERTY name:test2 visibility:public modality:FINAL [val]
     FIELD PROPERTY_BACKING_FIELD name:test2 type:kotlin.Function2<kotlin.Any, kotlin.Any, kotlin.Any> visibility:private [final,static]
       EXPRESSION_BODY
-        FUN_EXPR type=kotlin.Function1<kotlin.Any, IrErrorType> origin=LAMBDA
-          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:IrErrorType
+        FUN_EXPR type=kotlin.Function2<kotlin.Any, kotlin.Any, kotlin.Int> origin=LAMBDA
+          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.Any) returnType:kotlin.Int
+            VALUE_PARAMETER name:it index:0 type:kotlin.Any
             BLOCK_BODY
-              ERROR_CALL 'Unresolved reference: <Unresolved name: hashCode>#' type=IrErrorType
+              CALL 'public open fun hashCode (): kotlin.Int declared in kotlin.Any' type=kotlin.Int origin=null
+                $this: GET_VAR 'it: kotlin.Any declared in <root>.test2.<anonymous>' type=kotlin.Any origin=null
     FUN DEFAULT_PROPERTY_ACCESSOR name:<get-test2> visibility:public modality:FINAL <> () returnType:kotlin.Function2<kotlin.Any, kotlin.Any, kotlin.Any>
       correspondingProperty: PROPERTY name:test2 visibility:public modality:FINAL [val]
       BLOCK_BODY

compiler/testData/ir/irText/expressions/sam/samByProjectedType.fir.txt

@@ -1,9 +1,9 @@
 FILE fqName:<root> fileName:/samByProjectedType.kt
   FUN name:test1 visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
-      CALL 'public open fun bar (j: <root>.J<*>?): kotlin.Unit declared in <root>.H' type=kotlin.Unit origin=null
-        j: FUN_EXPR type=kotlin.Function1<kotlin.Any, kotlin.Any?> origin=LAMBDA
-          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (x:kotlin.Any) returnType:kotlin.Any?
+      ERROR_CALL 'Unresolved reference: <Inapplicable(INAPPLICABLE): [/H.bar]>#' type=IrErrorType
+        FUN_EXPR type=IrErrorType origin=LAMBDA
+          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (x:kotlin.Any) returnType:IrErrorType
             VALUE_PARAMETER name:x index:0 type:kotlin.Any
             BLOCK_BODY
-              GET_VAR 'x: kotlin.Any declared in <root>.test1.<anonymous>' type=kotlin.Any origin=null
+              ERROR_CALL 'Unresolved reference: x#' type=IrErrorType