NI: support inferring lambda receiver type by explicit receiver of anonymous function (which is another postponed argument)

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

@@ -10778,6 +10778,11 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirOldFronte
                 public void testLambdasInTryCatch() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/lambdasInTryCatch.kt");
                 }
+
+                @TestMetadata("takingExtensibilityFromDeclarationOfAnonymousFunction.kt")
+                public void testTakingExtensibilityFromDeclarationOfAnonymousFunction() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/takingExtensibilityFromDeclarationOfAnonymousFunction.kt");
+                }
             }
         }
 

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/components/PostponedArgumentInputTypesResolver.kt

@@ -92,18 +92,23 @@ class PostponedArgumentInputTypesResolver(
             }
         }?.toSet()
 
-        val annotations = functionalTypesFromConstraints?.run {
+        // An extension function flag can only come from a declaration of anonymous function: `select({ this + it }, fun Int.(x: Int) = 10)`
-            Annotations.create(map { it.type.annotations }.flatten())
+        val (parameterTypesFromDeclarationOfRelatedLambdas, isThereExtensionFunctionAmongRelatedLambdas) =
-        }
-
-        val parameterTypesFromDeclarationOfRelatedLambdas =
             getDeclaredParametersFromRelatedLambdas(argument, postponedArguments, variableDependencyProvider)
 
+        val annotationsFromConstraints = functionalTypesFromConstraints?.run {
+            Annotations.create(map { it.type.annotations }.flatten())
+        } ?: Annotations.EMPTY
+
+        val annotations = if (isThereExtensionFunctionAmongRelatedLambdas) {
+            annotationsFromConstraints.withExtensionFunctionAnnotation(expectedType.builtIns)
+        } else annotationsFromConstraints
+
         return ParameterTypesInfo(
             parameterTypesFromDeclaration,
             parameterTypesFromDeclarationOfRelatedLambdas,
             parameterTypesFromConstraints,
-            annotations ?: Annotations.EMPTY,
+            annotations,
             isSuspend = !functionalTypesFromConstraints.isNullOrEmpty() && functionalTypesFromConstraints.any { it.type.isSuspendFunctionTypeOrSubtype },
             isNullable = !functionalTypesFromConstraints.isNullOrEmpty() && functionalTypesFromConstraints.all { it.type.isMarkedNullable }
         )
@@ -113,24 +118,26 @@ class PostponedArgumentInputTypesResolver(
         argument: PostponedAtomWithRevisableExpectedType,
         postponedArguments: List<PostponedAtomWithRevisableExpectedType>,
         dependencyProvider: TypeVariableDependencyInformationProvider
-    ): Set<List<UnwrappedType?>>? {
+    ): Pair<Set<List<UnwrappedType?>>?, Boolean> {
-        fun PostponedAtomWithRevisableExpectedType.getExpectedTypeConstructor() = expectedType?.typeConstructor()
-
         val parameterTypesFromDeclarationOfRelatedLambdas = postponedArguments
             .filterIsInstance<LambdaWithTypeVariableAsExpectedTypeAtom>()
             .filter { it.parameterTypesFromDeclaration != null && it != argument }
             .mapNotNull { anotherArgument ->
-                val argumentExpectedTypeConstructor = argument.getExpectedTypeConstructor() ?: return@mapNotNull null
+                val argumentExpectedTypeConstructor = argument.expectedType?.typeConstructor() ?: return@mapNotNull null
-                val anotherArgumentExpectedTypeConstructor = anotherArgument.getExpectedTypeConstructor() ?: return@mapNotNull null
+                val anotherArgumentExpectedTypeConstructor = anotherArgument.expectedType.typeConstructor()
                 val areTypeVariablesRelated = dependencyProvider.areVariablesDependentShallowly(
-                    argumentExpectedTypeConstructor,
+                    argumentExpectedTypeConstructor, anotherArgumentExpectedTypeConstructor
-                    anotherArgumentExpectedTypeConstructor
                 )
+                val anotherAtom = anotherArgument.atom
+                val isAnonymousExtensionFunction = anotherAtom is FunctionExpression && anotherAtom.receiverType != null
+                val parameterTypesFromDeclarationOfRelatedLambda = anotherArgument.parameterTypesFromDeclaration
 
-                if (areTypeVariablesRelated) anotherArgument.parameterTypesFromDeclaration else null
+                if (areTypeVariablesRelated && parameterTypesFromDeclarationOfRelatedLambda != null) {
+                    parameterTypesFromDeclarationOfRelatedLambda to isAnonymousExtensionFunction
+                } else null
             }
 
-        return parameterTypesFromDeclarationOfRelatedLambdas.toSet().takeIf { it.isNotEmpty() }
+        return parameterTypesFromDeclarationOfRelatedLambdas.run { map { it.first }.toSet() to any { it.second } }
     }
 
     private fun Context.createTypeVariableForReturnType(argument: PostponedAtomWithRevisableExpectedType): NewTypeVariable {

compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/basic.kt

@@ -62,8 +62,7 @@ class A5<K, Q>: Function2<K, Q, Float> {
     override fun invoke(p1: K, p2: Q): Float = 5f
 }
 
-
+fun main() {
-fun test1() {
     // Inferring lambda parameter types by other lambda explicit parameters; expected type is type variable
     select(id1 { x -> <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>x<!>.inv() }, id2 { x: Int -> })
     select(id1 { <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>it<!>.inv() }, id2 { x: Int -> })
@@ -220,4 +219,17 @@ fun test1() {
     <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<*, kotlin.String?>")!>select(id { x: Int -> null }, id { x: String -> "" })<!>
     <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function<kotlin.Int>")!>select(id { x: Int -> 10 }, id { x: Int, y: Number -> TODO() })<!>
     val x68: String.(String) -> String = select(id { x: String, y: String -> "10" }, id { x: String, y: String -> "TODO()" })
+
+    // Anonymous functions
+    val x69: (C) -> Unit = selectB({ it }, { }, id(fun (x) { <!DEBUG_INFO_EXPRESSION_TYPE("A")!>x<!> }))
+    select(id1(fun(it) { <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>it<!>.inv() }), id1 { x: Number -> TODO() }, id1(id2(::takeInt)))
+    select(id(fun (it) { <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>it<!> }), id(id<(Int) -> Unit> { x: Number -> Unit }))
+    select(id(fun (it) { <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int")!>it<!>.inv() }), id<(Int) -> Unit> { })
+    val x70: (Int) -> Unit = selectNumber(id(fun (it) { }), id {}, id {})
+    val x71: String.() -> Unit = select(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String.() -> kotlin.Unit")!>id(fun String.() { })<!>, <!DEBUG_INFO_EXPRESSION_TYPE("(kotlin.String) -> kotlin.Unit")!>id(fun(x: String) {})<!>)
+    val x72: String.() -> Unit = select(fun String.() { }, fun(x: String) {}) // must be error
+    select(id(fun String.(x: String) {}), id(fun(x: String, y: String) { }), fun (x, y) { <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>x<!>;<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>y<!> })
+    select(id<Int.(String) -> Unit>(<!TYPE_MISMATCH!>fun (x, y) {}<!>), { x: Int, y: String -> x }) // receiver of anonymous function must be specified explicitly
+    select(id<Int.(String) -> Unit>(fun Int.(y) {}), { x: Int, y: String -> x })
+    <!DEBUG_INFO_EXPRESSION_TYPE("(kotlin.Number) -> java.io.Serializable")!>select(A3(), fun (x) = "", { a -> <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Number")!>a<!> })<!>
 }

compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/basic.txt

@@ -3,6 +3,7 @@ package
 public fun </*0*/ K> id(/*0*/ x: K): K
 public fun </*0*/ K> id1(/*0*/ x: K): K
 public fun </*0*/ L> id2(/*0*/ x: L): L
+public fun main(): kotlin.Unit
 public fun </*0*/ T> select(/*0*/ vararg x: T /*kotlin.Array<out T>*/): T
 public fun </*0*/ T : (A) -> kotlin.Unit> selectA(/*0*/ arg1: T, /*1*/ arg2: T, /*2*/ arg3: T): T
 public fun </*0*/ T : (B) -> kotlin.Unit> selectB(/*0*/ arg1: T, /*1*/ arg2: T, /*2*/ arg3: T): T
@@ -17,7 +18,6 @@ public fun </*0*/ T> takeLambdas(/*0*/ vararg x: (T) -> kotlin.Unit /*kotlin.Arr
 public fun </*0*/ T : R, /*1*/ R : L, /*2*/ L> takeLambdasWithDirectlyDependentTypeParameters(/*0*/ x: (T) -> kotlin.Unit, /*1*/ y: (R) -> kotlin.Unit, /*2*/ z: (L) -> kotlin.Unit): (T) -> kotlin.Unit
 public fun </*0*/ T, /*1*/ R : T, /*2*/ L : R> takeLambdasWithInverselyDependentTypeParameters(/*0*/ x: (T) -> kotlin.Unit, /*1*/ y: (R) -> kotlin.Unit, /*2*/ z: (L) -> kotlin.Unit): (T) -> kotlin.Unit
 public fun takeNumber(/*0*/ i: kotlin.Number): kotlin.Unit
-public fun test1(): kotlin.Unit
 public fun withOverload(/*0*/ i: kotlin.Int): kotlin.Unit
 public fun withOverload(/*0*/ s: kotlin.String): kotlin.Unit
 

compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/takingExtensibilityFromDeclarationOfAnonymousFunction.fir.kt

@@ -0,0 +1,8 @@
+// !DIAGNOSTICS: -UNUSED_ANONYMOUS_PARAMETER
+
+fun <T> id(x: T) = x
+fun <T> select(vararg x: T) = x[0]
+
+val x1 = select(id { this }, fun Int.() = this)
+val x2 = select(id { this + <!UNRESOLVED_REFERENCE!>it<!>.<!UNRESOLVED_REFERENCE!>inv<!>() }, fun Int.(x: Int) = this)
+val x3 = select(id { this.<!UNRESOLVED_REFERENCE!>length<!> + <!UNRESOLVED_REFERENCE!>it<!>.<!UNRESOLVED_REFERENCE!>inv<!>() }, fun String.(x: Int) = length)

compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/takingExtensibilityFromDeclarationOfAnonymousFunction.kt

@@ -0,0 +1,8 @@
+// !DIAGNOSTICS: -UNUSED_ANONYMOUS_PARAMETER
+
+fun <T> id(x: T) = x
+fun <T> select(vararg x: T) = x[0]
+
+val x1 = select(id { this }, fun Int.() = this)
+val x2 = select(id { this + it.inv() }, fun Int.(x: Int) = this)
+val x3 = select(id { this.length + it.inv() }, fun String.(x: Int) = length)

compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/takingExtensibilityFromDeclarationOfAnonymousFunction.txt

@@ -0,0 +1,7 @@
+package
+
+public val x1: kotlin.Int.() -> kotlin.Int
+public val x2: kotlin.Int.(kotlin.Int) -> kotlin.Int
+public val x3: kotlin.String.(kotlin.Int) -> kotlin.Int
+public fun </*0*/ T> id(/*0*/ x: T): T
+public fun </*0*/ T> select(/*0*/ vararg x: T /*kotlin.Array<out T>*/): T

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -10785,6 +10785,11 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTestWithFirVali
                     public void testLambdasInTryCatch() throws Exception {
                         runTest("compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/lambdasInTryCatch.kt");
                     }
+
+                    @TestMetadata("takingExtensibilityFromDeclarationOfAnonymousFunction.kt")
+                    public void testTakingExtensibilityFromDeclarationOfAnonymousFunction() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/takingExtensibilityFromDeclarationOfAnonymousFunction.kt");
+                    }
                 }
             }
 

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsUsingJavacTestGenerated.java

@@ -10780,6 +10780,11 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                     public void testLambdasInTryCatch() throws Exception {
                         runTest("compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/lambdasInTryCatch.kt");
                     }
+
+                    @TestMetadata("takingExtensibilityFromDeclarationOfAnonymousFunction.kt")
+                    public void testTakingExtensibilityFromDeclarationOfAnonymousFunction() throws Exception {
+                        runTest("compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/takingExtensibilityFromDeclarationOfAnonymousFunction.kt");
+                    }
                 }
             }