NI: Analyze lambdas which are return arguments of another lambda

^KT-36044 Fixed

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirOldFrontendDiagnosticsTestGenerated.java

@@ -10106,6 +10106,11 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirOldFronte
             runTest("compiler/testData/diagnostics/tests/inference/kt35702.kt");
         }
 
+        @TestMetadata("kt36044.kt")
+        public void testKt36044() throws Exception {
+            runTest("compiler/testData/diagnostics/tests/inference/kt36044.kt");
+        }
+
         @TestMetadata("kt36819.kt")
         public void testKt36819() throws Exception {
             runTest("compiler/testData/diagnostics/tests/inference/kt36819.kt");

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

@@ -150,7 +150,10 @@ class PostponedArgumentsAnalyzer(
         val subResolvedKtPrimitives = allReturnArguments.map {
             resolveKtPrimitive(
                 c.getBuilder(), it, lambda.returnType.let(::substitute), diagnosticHolder, ReceiverInfo.notReceiver, convertedType = null
-            )
+            ).apply {
+                if (this is LambdaWithTypeVariableAsExpectedTypeAtom)
+                    isReturnArgumentOfAnotherLambda = true
+            }
         }
 
         if (!returnArgumentsInfo.returnArgumentsExist) {

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

@@ -5,11 +5,13 @@
 
 package org.jetbrains.kotlin.resolve.calls.inference.components
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
 import org.jetbrains.kotlin.builtins.isBuiltinFunctionalTypeOrSubtype
 import org.jetbrains.kotlin.builtins.isExtensionFunctionType
 import org.jetbrains.kotlin.resolve.calls.components.transformToResolvedLambda
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
+import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.types.*
@@ -103,12 +105,12 @@ class KotlinConstraintSystemCompleter(
 
             if (
                 completionMode == ConstraintSystemCompletionMode.FULL &&
-                resolveLambdaOrCallableReferenceWithTypeVariableAsExpectedType(
+                c.resolveLambdaByAdditionalConditions(
-                    c,
                     variableForFixation,
                     topLevelAtoms,
                     diagnosticsHolder,
-                    analyze
+                    analyze,
+                    variableFixationFinder
                 )
             ) {
                 continue
@@ -141,15 +143,67 @@ class KotlinConstraintSystemCompleter(
     /*
      * returns true -> analyzed
      */
-    private fun resolveLambdaOrCallableReferenceWithTypeVariableAsExpectedType(
+    private fun Context.resolveLambdaByAdditionalConditions(
-        c: Context,
         variableForFixation: VariableFixationFinder.VariableForFixation,
         topLevelAtoms: List<ResolvedAtom>,
         diagnosticsHolder: KotlinDiagnosticsHolder,
+        analyze: (PostponedResolvedAtom) -> Unit,
+        fixationFinder: VariableFixationFinder
+    ): Boolean {
+        val postponedArguments = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
+
+        return resolveLambdaOrCallableReferenceWithTypeVariableAsExpectedType(
+            variableForFixation,
+            postponedArguments,
+            diagnosticsHolder,
+            analyze
+        ) || resolveLambdaWhichIsReturnArgument(postponedArguments, diagnosticsHolder, analyze, fixationFinder)
+    }
+
+    /*
+     * returns true -> analyzed
+     */
+    private fun Context.resolveLambdaWhichIsReturnArgument(
+        postponedArguments: List<PostponedResolvedAtom>,
+        diagnosticsHolder: KotlinDiagnosticsHolder,
+        analyze: (PostponedResolvedAtom) -> Unit,
+        fixationFinder: VariableFixationFinder
+    ): Boolean {
+        if (this !is NewConstraintSystem) return false
+
+        val isReturnArgumentOfAnotherLambda = postponedArguments.any {
+            it is LambdaWithTypeVariableAsExpectedTypeAtom && it.isReturnArgumentOfAnotherLambda
+        }
+        val postponedAtom = postponedArguments.firstOrNull() ?: return false
+        val atomExpectedType = postponedAtom.expectedType
+
+        val shouldAnalyzeByPresenceLambdaAsReturnArgument =
+            isReturnArgumentOfAnotherLambda && atomExpectedType != null &&
+                    with(fixationFinder) { variableHasTrivialOrNonProperConstraints(atomExpectedType.constructor) }
+
+        if (!shouldAnalyzeByPresenceLambdaAsReturnArgument)
+            return false
+
+        val expectedTypeVariable =
+            atomExpectedType?.constructor?.takeIf { it in this.getBuilder().currentStorage().allTypeVariables } ?: return false
+
+        analyze(preparePostponedAtom(expectedTypeVariable, postponedAtom, expectedTypeVariable.builtIns, diagnosticsHolder) ?: return false)
+
+        return true
+    }
+
+    /*
+     * returns true -> analyzed
+     */
+    private fun Context.resolveLambdaOrCallableReferenceWithTypeVariableAsExpectedType(
+        variableForFixation: VariableFixationFinder.VariableForFixation,
+        postponedArguments: List<PostponedResolvedAtom>,
+        diagnosticsHolder: KotlinDiagnosticsHolder,
         analyze: (PostponedResolvedAtom) -> Unit
     ): Boolean {
-        val variable = variableForFixation.variable as TypeConstructor
+        if (this !is NewConstraintSystem) return false
-        val postponedArguments = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
+
+        val variable = variableForFixation.variable as? TypeConstructor ?: return false
         val hasProperAtom = postponedArguments.any {
             when (it) {
                 is LambdaWithTypeVariableAsExpectedTypeAtom,
@@ -158,21 +212,33 @@ class KotlinConstraintSystemCompleter(
             }
         }
 
-        if (
-            !hasProperAtom &&
-            variableForFixation.hasProperConstraint &&
-            !variableForFixation.hasOnlyTrivialProperConstraint
-        ) return false
-
         val postponedAtom = postponedArguments.firstOrNull() ?: return false
+        val expectedTypeAtom = postponedAtom.expectedType
+        val expectedTypeVariable =
+            expectedTypeAtom?.constructor?.takeIf { it in this.getBuilder().currentStorage().allTypeVariables } ?: variable
 
-        val builtIns = (variable as TypeVariableTypeConstructor).builtIns
+        val shouldAnalyzeByEqualityExpectedTypeToVariable =
-        val csBuilder = (c as NewConstraintSystemImpl).getBuilder()
+            hasProperAtom || !variableForFixation.hasProperConstraint || variableForFixation.hasOnlyTrivialProperConstraint
 
-        val expectedTypeVariable = postponedAtom.expectedType?.constructor?.takeIf { it in c.allTypeVariables } ?: variable
+        if (!shouldAnalyzeByEqualityExpectedTypeToVariable)
-        val atomToAnalyze = when (postponedAtom) {
+            return false
+
+        analyze(preparePostponedAtom(expectedTypeVariable, postponedAtom, variable.builtIns, diagnosticsHolder) ?: return false)
+
+        return true
+    }
+
+    private fun Context.preparePostponedAtom(
+        expectedTypeVariable: TypeConstructor,
+        postponedAtom: PostponedResolvedAtom,
+        builtIns: KotlinBuiltIns,
+        diagnosticsHolder: KotlinDiagnosticsHolder
+    ): PostponedResolvedAtom? {
+        val csBuilder = (this as? NewConstraintSystem)?.getBuilder() ?: return null
+
+        return when (postponedAtom) {
             is PostponedCallableReferenceAtom -> postponedAtom.preparePostponedAtomWithTypeVariableAsExpectedType(
-                c, csBuilder, expectedTypeVariable,
+                this, csBuilder, expectedTypeVariable,
                 parameterTypes = null,
                 isSuitable = KotlinType::isBuiltinFunctionalTypeOrSubtype,
                 typeVariableCreator = { TypeVariableForCallableReferenceReturnType(builtIns, "_Q") },
@@ -183,7 +249,7 @@ class KotlinConstraintSystemCompleter(
                 }
             )
             is LambdaWithTypeVariableAsExpectedTypeAtom -> postponedAtom.preparePostponedAtomWithTypeVariableAsExpectedType(
-                c, csBuilder, expectedTypeVariable,
+                this, csBuilder, expectedTypeVariable,
                 parameterTypes = postponedAtom.atom.parametersTypes,
                 isSuitable = KotlinType::isBuiltinFunctionalType,
                 typeVariableCreator = { TypeVariableForLambdaReturnType(postponedAtom.atom, builtIns, "_R") },
@@ -191,10 +257,8 @@ class KotlinConstraintSystemCompleter(
                     postponedAtom.transformToResolvedLambda(csBuilder, diagnosticsHolder, expectedType, returnVariable)
                 }
             )
-            else -> return false
+            else -> null
         }
-        analyze(atomToAnalyze)
-        return true
     }
 
     private inline fun <T : PostponedResolvedAtom, V : NewTypeVariable> T.preparePostponedAtomWithTypeVariableAsExpectedType(

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

@@ -86,6 +86,13 @@ class VariableFixationFinder(
         }
     }
 
+    fun Context.variableHasTrivialOrNonProperConstraints(variable: TypeConstructorMarker): Boolean {
+        return notFixedTypeVariables[variable]?.constraints?.all { constraint ->
+            val isProperConstraint = isProperArgumentConstraint(constraint)
+            isProperConstraint && trivialConstraintTypeInferenceOracle.isNotInterestingConstraint(constraint) || !isProperConstraint
+        } ?: false
+    }
+
     private fun Context.findTypeVariableForFixation(
         allTypeVariables: List<TypeConstructorMarker>,
         postponedArguments: List<PostponedResolvedAtomMarker>,
@@ -120,13 +127,6 @@ class VariableFixationFinder(
         return false
     }
 
-    private fun Context.variableHasTrivialOrNonProperConstraints(variable: TypeConstructorMarker): Boolean {
-        return notFixedTypeVariables[variable]?.constraints?.all { constraint ->
-            val isProperConstraint = isProperArgumentConstraint(constraint)
-            isProperConstraint && trivialConstraintTypeInferenceOracle.isNotInterestingConstraint(constraint) || !isProperConstraint
-        } ?: false
-    }
-
     private fun Context.variableHasProperArgumentConstraints(variable: TypeConstructorMarker): Boolean =
         notFixedTypeVariables[variable]?.constraints?.any { isProperArgumentConstraint(it) } ?: false
 

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

@@ -109,6 +109,7 @@ class LambdaWithTypeVariableAsExpectedTypeAtom(
 ) : PostponedResolvedAtom() {
     override val inputTypes: Collection<UnwrappedType> get() = listOf(expectedType)
     override val outputType: UnwrappedType? get() = null
+    var isReturnArgumentOfAnotherLambda: Boolean = false
 
     fun setAnalyzed(resolvedLambdaAtom: ResolvedLambdaAtom) {
         setAnalyzedResults(listOf(resolvedLambdaAtom))

compiler/testData/diagnostics/tests/callableReference/generic/nestedCallWithOverload.fir.fail

@@ -0,0 +1,8 @@
+Failures detected in FirBodyResolveTransformerAdapter, file: /nestedCallWithOverload.fir.kt
+Cause: java.lang.RuntimeException: While resolving call R?C|/baz|(R?C|/id|(<L> = id@fun <implicit>.<anonymous>(): <implicit> {
+    it#.inv#()
+}
+), R|/id|<R|(kotlin/Int) -> kotlin/Unit|>(<L> = id@fun <anonymous>(it: R|kotlin/Int|): R|kotlin/Unit| {
+    Unit
+}
+))

compiler/testData/diagnostics/tests/callableReference/generic/nestedCallWithOverload.fir.kt

@@ -19,4 +19,6 @@ fun test() {
     baz<(Int) -> Unit>(id(::foo), id(id(::foo)))
     baz(id(::foo), id(id<(Int) -> Unit>(::foo)))
     baz(id(<!UNRESOLVED_REFERENCE!>::foo<!>), id<(Int) -> Unit>(id(::foo)))
+
+    baz(id { it.inv() }, id<(Int) -> Unit> { })
 }

compiler/testData/diagnostics/tests/callableReference/generic/nestedCallWithOverload.kt

@@ -19,4 +19,6 @@ fun test() {
     baz<(Int) -> Unit>(id(::foo), id(id(::foo)))
     baz(id(::foo), id(id<(Int) -> Unit>(::foo)))
     baz(id(::foo), id<(Int) -> Unit>(id(::foo)))
+
+    baz(id { it.inv() }, id<(Int) -> Unit> { })
 }

compiler/testData/diagnostics/tests/inference/kt36044.fir.fail

@@ -0,0 +1,8 @@
+Failures detected in FirBodyResolveTransformerAdapter, file: /kt36044.fir.kt
+Cause: java.lang.RuntimeException: While resolving call R?C|/select|(Char(a), R?C|/map|(<L> = map@fun <anonymous>(): <ERROR TYPE REF: No type for block> {
+    ^ map@fun <implicit>.<anonymous>(): <implicit> {
+        String()
+    }
+
+}
+))

compiler/testData/diagnostics/tests/inference/kt36044.fir.kt

@@ -0,0 +1,9 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// Issue: KT-36044
+
+fun <A> select(x: A, f: () -> A) = f()
+fun <B> map(f: () -> B) = f()
+
+fun main() {
+    select('a', map { { "" } })
+}

compiler/testData/diagnostics/tests/inference/kt36044.kt

@@ -0,0 +1,9 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// Issue: KT-36044
+
+fun <A> select(x: A, f: () -> A) = f()
+fun <B> map(f: () -> B) = f()
+
+fun main() {
+    select('a', map { { "" } })
+}

compiler/testData/diagnostics/tests/inference/kt36044.txt

@@ -0,0 +1,5 @@
+package
+
+public fun main(): kotlin.Unit
+public fun </*0*/ B> map(/*0*/ f: () -> B): B
+public fun </*0*/ A> select(/*0*/ x: A, /*1*/ f: () -> A): A

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

@@ -10113,6 +10113,11 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTestWithFirVali
                 runTest("compiler/testData/diagnostics/tests/inference/kt35702.kt");
             }
 
+            @TestMetadata("kt36044.kt")
+            public void testKt36044() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/kt36044.kt");
+            }
+
             @TestMetadata("kt36819.kt")
             public void testKt36819() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/inference/kt36819.kt");

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

@@ -10108,6 +10108,11 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                 runTest("compiler/testData/diagnostics/tests/inference/kt35702.kt");
             }
 
+            @TestMetadata("kt36044.kt")
+            public void testKt36044() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/kt36044.kt");
+            }
+
             @TestMetadata("kt36819.kt")
             public void testKt36819() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/inference/kt36819.kt");