FIR: Order type variables for fixation

Otherwise there are subtle issues with captured types
(see FirDiagnosticsSmokeTestGenerated$Inference.testDependantOnVariance)

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

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.fir.resolve.calls
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
@@ -18,6 +17,7 @@ import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeVariable
 import org.jetbrains.kotlin.fir.types.FirTypeProjection
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
@@ -79,6 +79,7 @@ class Candidate(
     val samResolver get() = bodyResolveComponents.samResolver
 
     lateinit var substitutor: ConeSubstitutor
+    lateinit var freshVariables: List<ConeTypeVariable>
     var resultingTypeForCallableReference: ConeKotlinType? = null
     var outerConstraintBuilderEffect: (ConstraintSystemOperation.() -> Unit)? = null
 

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

@@ -25,12 +25,13 @@ internal object CreateFreshTypeVariableSubstitutorStage : ResolutionStage() {
         val declaration = candidate.symbol.fir
         if (declaration !is FirTypeParametersOwner || declaration.typeParameters.isEmpty()) {
             candidate.substitutor = ConeSubstitutor.Empty
+            candidate.freshVariables = emptyList()
             return
         }
         val csBuilder = candidate.system.getBuilder()
         val (substitutor, freshVariables) = createToFreshVariableSubstitutorAndAddInitialConstraints(declaration, candidate, csBuilder)
         candidate.substitutor = substitutor
-
+        candidate.freshVariables = freshVariables
 
         // bad function -- error on declaration side
         if (csBuilder.hasContradiction) {

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

@@ -6,7 +6,9 @@
 package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.returnExpressions
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeVariable
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
 import org.jetbrains.kotlin.resolve.calls.inference.components.TypeVariableDirectionCalculator
@@ -15,6 +17,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraint
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtom
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
+import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.isIntegerLiteralTypeConstructor
 import org.jetbrains.kotlin.types.model.typeConstructor
 import org.jetbrains.kotlin.utils.addIfNotNull
@@ -88,7 +91,7 @@ private fun Candidate.hasProperNonTrivialLowerConstraints(components: InferenceC
 
 
 class ConstraintSystemCompleter(val components: InferenceComponents) {
-    val variableFixationFinder = VariableFixationFinder(components.trivialConstraintTypeInferenceOracle)
+    private val variableFixationFinder = VariableFixationFinder(components.trivialConstraintTypeInferenceOracle)
     fun complete(
         c: KotlinConstraintSystemCompleter.Context,
         completionMode: KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode,
@@ -100,9 +103,7 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
         while (true) {
             if (analyzePostponeArgumentIfPossible(c, topLevelAtoms, analyze)) continue
 
-
+            val allTypeVariables = getOrderedAllTypeVariables(c, topLevelAtoms)
-//            val allTypeVariables = getOrderedAllTypeVariables(c, collectVariablesFromContext, topLevelAtoms)
-            val allTypeVariables = c.notFixedTypeVariables.keys.toList()
             val postponedKtPrimitives = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
             val variableForFixation =
                 variableFixationFinder.findFirstVariableForFixation(
@@ -138,6 +139,43 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
         }
     }
 
+    private fun getOrderedAllTypeVariables(
+        c: KotlinConstraintSystemCompleter.Context,
+        topLevelAtoms: List<FirStatement>
+    ): List<TypeConstructorMarker> {
+        val result = LinkedHashSet<TypeConstructorMarker>(c.notFixedTypeVariables.size)
+        fun ConeTypeVariable?.toTypeConstructor(): TypeConstructorMarker? =
+            this?.typeConstructor?.takeIf { it in c.notFixedTypeVariables.keys }
+
+        fun FirStatement.collectAllTypeVariables() {
+            this.processAllContainingCallCandidates(processBlocks = true) { candidate ->
+                candidate.freshVariables.mapNotNullTo(result) { typeVariable ->
+                    typeVariable.toTypeConstructor()
+                }
+
+                for (lambdaAtom in candidate.postponedAtoms.filterIsInstance<ResolvedLambdaAtom>()) {
+                    result.addIfNotNull(lambdaAtom.typeVariableForLambdaReturnType.toTypeConstructor())
+                    if (lambdaAtom.analyzed) {
+                        for (firStatement in lambdaAtom.returnStatements) {
+                            firStatement.collectAllTypeVariables()
+                        }
+                    }
+                }
+            }
+        }
+
+        for (topLevel in topLevelAtoms) {
+            topLevel.collectAllTypeVariables()
+        }
+
+        require(result.size == c.notFixedTypeVariables.size) {
+            val notFoundTypeVariables = c.notFixedTypeVariables.keys.toMutableSet().apply { removeAll(result) }
+            "Not all type variables found: $notFoundTypeVariables"
+        }
+
+        return result.toList()
+    }
+
     private fun fixVariable(
         c: KotlinConstraintSystemCompleter.Context,
         topLevelType: KotlinTypeMarker,
@@ -172,48 +210,62 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
     }
 
     private fun getOrderedNotAnalyzedPostponedArguments(topLevelAtoms: List<FirStatement>): List<PostponedResolvedAtomMarker> {
-        fun FirStatement.process(to: MutableList<PostponedResolvedAtomMarker>) {
-            when (this) {
-                is FirFunctionCall -> {
-                    val candidate = (this.calleeReference as? FirNamedReferenceWithCandidate)?.candidate
-                    candidate?.postponedAtoms?.forEach {
-                        to.addIfNotNull(it.safeAs<PostponedResolvedAtomMarker>()?.takeUnless { it.analyzed })
-                    }
-                    this.arguments.forEach { it.process(to) }
-                }
-                is FirWhenExpression -> {
-                    val candidate = (this.calleeReference as? FirNamedReferenceWithCandidate)?.candidate
-                    candidate?.postponedAtoms?.forEach {
-                        to.addIfNotNull(it.safeAs<PostponedResolvedAtomMarker>()?.takeUnless { it.analyzed })
-                    }
-                    this.branches.forEach { it.result.process(to) }
-                }
-
-                is FirTryExpression -> {
-                    val candidate = (this.calleeReference as? FirNamedReferenceWithCandidate)?.candidate
-                    candidate?.postponedAtoms?.forEach {
-                        to.addIfNotNull(it.safeAs<PostponedResolvedAtomMarker>()?.takeUnless { it.analyzed })
-                    }
-                    tryBlock.process(to)
-                    catches.forEach { it.block.process(to) }
-                }
-
-                is FirWrappedArgumentExpression -> this.expression.process(to)
-                // TOOD: WTF?
-            }
-//            if (analyzed) {
-//                subResolvedAtoms.forEach { it.process(to) }
-//            }
-        }
-
         val notAnalyzedArguments = arrayListOf<PostponedResolvedAtomMarker>()
         for (primitive in topLevelAtoms) {
-            primitive.process(notAnalyzedArguments)
+            primitive.processAllContainingCallCandidates(
+                // TODO: remove this argument and relevant parameter
+                // Currently, it's used because otherwise problem happens with a lambda in a try-block (see tryWithLambdaInside test)
+                processBlocks = false
+            ) { candidate ->
+                candidate.postponedAtoms.forEach {
+                    notAnalyzedArguments.addIfNotNull(it.safeAs<PostponedResolvedAtomMarker>()?.takeUnless { it.analyzed })
+                }
+            }
         }
 
         return notAnalyzedArguments
     }
 
+    private fun FirStatement.processAllContainingCallCandidates(processBlocks: Boolean, processor: (Candidate) -> Unit) {
+        when (this) {
+            is FirFunctionCall -> {
+                processCandidateIfApplicable(processor)
+                this.arguments.forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
+            }
+
+            is FirWhenExpression -> {
+                processCandidateIfApplicable(processor)
+                this.branches.forEach { it.result.processAllContainingCallCandidates(processBlocks, processor) }
+            }
+
+            is FirTryExpression -> {
+                processCandidateIfApplicable(processor)
+                tryBlock.processAllContainingCallCandidates(processBlocks, processor)
+                catches.forEach { it.block.processAllContainingCallCandidates(processBlocks, processor) }
+            }
+
+            is FirQualifiedAccessExpression -> {
+                processCandidateIfApplicable(processor)
+            }
+
+            is FirVariableAssignment -> {
+                processCandidateIfApplicable(processor)
+                rValue.processAllContainingCallCandidates(processBlocks, processor)
+            }
+
+            is FirWrappedArgumentExpression -> this.expression.processAllContainingCallCandidates(processBlocks, processor)
+            is FirBlock -> {
+                if (processBlocks) {
+                    this.returnExpressions().forEach { it.processAllContainingCallCandidates(processBlocks, processor) }
+                }
+            }
+        }
+    }
+
+    private inline fun FirResolvable.processCandidateIfApplicable(processor: (Candidate) -> Unit) {
+        (calleeReference as? FirNamedReferenceWithCandidate)?.candidate?.let(processor)
+    }
+
     private fun canWeAnalyzeIt(c: KotlinConstraintSystemCompleter.Context, argument: PostponedResolvedAtomMarker): Boolean {
         if (argument.analyzed) return false
 

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
+import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.DoubleColonLHS
 import org.jetbrains.kotlin.fir.resolve.constructType
 import org.jetbrains.kotlin.fir.resolve.createFunctionalType
@@ -199,6 +200,7 @@ class ResolvedLambdaAtom(
 ) : PostponedResolvedAtomMarker {
 
     override var analyzed: Boolean = false
+    lateinit var returnStatements: List<FirStatement>
 
 //    lateinit var resultArguments: List<KotlinCallArgument>
 //        private set

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

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.FirCallResolver
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
-import org.jetbrains.kotlin.fir.diagnostics.FirSimpleDiagnostic
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.references.impl.FirErrorNamedReferenceImpl
@@ -177,6 +176,7 @@ class PostponedArgumentsAnalyzer(
         }
 
         lambda.analyzed = true
+        lambda.returnStatements = returnArguments
         //lambda.setAnalyzedResults(returnArguments, subResolvedKtPrimitives)
 
 //        if (inferenceSession != null) {

compiler/fir/resolve/testData/resolve/stdlib/problems/tryWithLambdaInside.kt

@@ -0,0 +1,15 @@
+// FULL_JDK
+
+fun foo(): List<String> = TODO()
+
+fun <T> ba(): List<T> = TODO()
+
+fun bar() =
+    try {
+        foo().filter {
+            // Lambda remains effectively unresolved
+            it.length > 2
+        }
+    } finally {
+        ba()
+    }

compiler/fir/resolve/testData/resolve/stdlib/problems/tryWithLambdaInside.txt

@@ -0,0 +1,19 @@
+FILE: tryWithLambdaInside.kt
+    public final fun foo(): R|kotlin/collections/List<kotlin/String>| {
+        ^foo R|kotlin/TODO|()
+    }
+    public final fun <T> ba(): R|kotlin/collections/List<T>| {
+        ^ba R|kotlin/TODO|()
+    }
+    public final fun bar(): R|kotlin/collections/List<kotlin/String>| {
+        ^bar try {
+            R|/foo|().R|kotlin/collections/filter|<R|kotlin/String|>(<L> = filter@fun <implicit>.<anonymous>(): <implicit> <kind=UNKNOWN>  {
+                >(it#.length#, Int(2))
+            }
+            )
+        }
+        finally {
+            R|/ba|<R|kotlin/Nothing|>()
+        }
+
+    }

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

@@ -576,5 +576,10 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
         public void testCloneArray() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/stdlib/problems/cloneArray.kt");
         }
+
+        @TestMetadata("tryWithLambdaInside.kt")
+        public void testTryWithLambdaInside() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/stdlib/problems/tryWithLambdaInside.kt");
+        }
     }
 }