FIR: Force processing postponed arguments nested in lambdas

See the call-site of processAllContainingCallCandidates in getOrderedNotAnalyzedPostponedArguments

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

@@ -6,8 +6,8 @@
 package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.types.ConeIntegerLiteralType
 import org.jetbrains.kotlin.fir.returnExpressions
+import org.jetbrains.kotlin.fir.types.ConeIntegerLiteralType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeTypeVariable
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -157,11 +157,6 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
 
                 for (lambdaAtom in candidate.postponedAtoms.filterIsInstance<ResolvedLambdaAtom>()) {
                     result.addIfNotNull(lambdaAtom.typeVariableForLambdaReturnType.toTypeConstructor())
-                    if (lambdaAtom.analyzed) {
-                        for (firStatement in lambdaAtom.returnStatements) {
-                            firStatement.collectAllTypeVariables()
-                        }
-                    }
                 }
             }
         }
@@ -231,32 +226,32 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
     private fun FirStatement.processAllContainingCallCandidates(processBlocks: Boolean, processor: (Candidate) -> Unit) {
         when (this) {
             is FirFunctionCall -> {
-                processCandidateIfApplicable(processor)
+                processCandidateIfApplicable(processor, processBlocks)
                 this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
             }
 
             is FirWhenExpression -> {
-                processCandidateIfApplicable(processor)
+                processCandidateIfApplicable(processor, processBlocks)
                 this.branches.forEach { it.result.processAllContainingCallCandidates(processBlocks, processor) }
             }
 
             is FirTryExpression -> {
-                processCandidateIfApplicable(processor)
+                processCandidateIfApplicable(processor, processBlocks)
                 tryBlock.processAllContainingCallCandidates(processBlocks, processor)
                 catches.forEach { it.block.processAllContainingCallCandidates(processBlocks, processor) }
             }
 
             is FirCheckNotNullCall -> {
-                processCandidateIfApplicable(processor)
+                processCandidateIfApplicable(processor, processBlocks)
                 this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
             }
 
             is FirQualifiedAccessExpression -> {
-                processCandidateIfApplicable(processor)
+                processCandidateIfApplicable(processor, processBlocks)
             }
 
             is FirVariableAssignment -> {
-                processCandidateIfApplicable(processor)
+                processCandidateIfApplicable(processor, processBlocks)
                 rValue.processAllContainingCallCandidates(processBlocks, processor)
             }
 
@@ -268,14 +263,26 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
             }
 
             is FirDelegatedConstructorCall -> {
-                processCandidateIfApplicable(processor)
+                processCandidateIfApplicable(processor, processBlocks)
                 this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
             }
         }
     }
 
-    private inline fun FirResolvable.processCandidateIfApplicable(processor: (Candidate) -> Unit) {
+    private fun FirResolvable.processCandidateIfApplicable(
-        (calleeReference as? FirNamedReferenceWithCandidate)?.candidate?.let(processor)
+        processor: (Candidate) -> Unit,
+        processBlocks: Boolean
+    ) {
+        val candidate = (calleeReference as? FirNamedReferenceWithCandidate)?.candidate ?: return
+        processor(candidate)
+
+        for (atom in candidate.postponedAtoms) {
+            if (atom !is ResolvedLambdaAtom || !atom.analyzed) continue
+
+            atom.returnStatements.forEach {
+                it.processAllContainingCallCandidates(processBlocks, processor)
+            }
+        }
     }
 
     private fun canWeAnalyzeIt(c: KotlinConstraintSystemCompleter.Context, argument: PostponedResolvedAtomMarker): Boolean {

compiler/fir/resolve/testData/resolve/inference/nestedLambdas.kt

@@ -0,0 +1,21 @@
+fun <T> myRun(computable: () -> T): T = TODO()
+
+interface Inv<W>
+interface MyMap<K, V> {
+    val k: K
+    val v: V
+}
+
+val w: Inv<String> = TODO()
+
+public fun <X, K> Inv<X>.associateBy1(keySelector: (X) -> K): MyMap<K, String> = TODO()
+
+val x = myRun {
+    w.associateBy1 { f -> f.length }
+}
+
+fun foo(m: MyMap<Int, String>) {}
+
+fun main() {
+    foo(x)
+}

compiler/fir/resolve/testData/resolve/inference/nestedLambdas.txt

@@ -0,0 +1,32 @@
+FILE: nestedLambdas.kt
+    public final fun <T> myRun(computable: R|() -> T|): R|T| {
+        ^myRun R|kotlin/TODO|()
+    }
+    public abstract interface Inv<W> : R|kotlin/Any| {
+    }
+    public abstract interface MyMap<K, V> : R|kotlin/Any| {
+        public abstract val k: R|K|
+            public get(): R|K|
+
+        public abstract val v: R|V|
+            public get(): R|V|
+
+    }
+    public final val w: R|Inv<kotlin/String>| = R|kotlin/TODO|()
+        public get(): R|Inv<kotlin/String>|
+    public final fun <X, K> R|Inv<X>|.associateBy1(keySelector: R|(X) -> K|): R|MyMap<K, kotlin/String>| {
+        ^associateBy1 R|kotlin/TODO|()
+    }
+    public final val x: R|MyMap<kotlin/Int, kotlin/String>| = R|/myRun|<R|MyMap<kotlin/Int, kotlin/String>|>(<L> = myRun@fun <anonymous>(): R|MyMap<kotlin/Int, kotlin/String>| {
+        R|/w|.R|/associateBy1|<R|kotlin/String|, R|kotlin/Int|>(<L> = associateBy1@fun <anonymous>(f: R|kotlin/String|): R|kotlin/Int| {
+            R|<local>/f|.R|kotlin/String.length|
+        }
+        )
+    }
+    )
+        public get(): R|MyMap<kotlin/Int, kotlin/String>|
+    public final fun foo(m: R|MyMap<kotlin/Int, kotlin/String>|): R|kotlin/Unit| {
+    }
+    public final fun main(): R|kotlin/Unit| {
+        R|/foo|(R|/x|)
+    }

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

@@ -904,6 +904,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/resolve/testData/resolve/inference"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
         }
 
+        @TestMetadata("nestedLambdas.kt")
+        public void testNestedLambdas() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/inference/nestedLambdas.kt");
+        }
+
         @TestMetadata("receiverWithCapturedType.kt")
         public void testReceiverWithCapturedType() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/inference/receiverWithCapturedType.kt");

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

@@ -904,6 +904,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/resolve/testData/resolve/inference"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
         }
 
+        @TestMetadata("nestedLambdas.kt")
+        public void testNestedLambdas() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/inference/nestedLambdas.kt");
+        }
+
         @TestMetadata("receiverWithCapturedType.kt")
         public void testReceiverWithCapturedType() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/inference/receiverWithCapturedType.kt");

compiler/testData/diagnostics/tests/delegatedProperty/inference/delegateExpressionAsLambda.fir.kt

@@ -7,7 +7,7 @@ fun test(i: Int) {
         createSample(i) { it.toString() }
     }
 
-    <!INAPPLICABLE_CANDIDATE!>takeSample<!>(bad)
+    takeSample(bad)
 }
 
 fun <T> myLazyDelegate(i: () -> T): LazyDelegate<T> = LazyDelegate(i())
@@ -19,4 +19,4 @@ class LazyDelegate<T>(val v: T) {
 class Sample<K, V>
 
 fun takeSample(g: Sample<Int, String>) {}
-fun <T, S> createSample(i: T, a: (T) -> S): Sample<T, S> = TODO()
+fun <T, S> createSample(i: T, a: (T) -> S): Sample<T, S> = TODO()