K2/PCLA: add explicit withPCLASession: Boolean parameter to analyzeLambda

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirOverloadByLambdaReturnTypeResolver.kt

@@ -118,6 +118,7 @@ class FirOverloadByLambdaReturnTypeResolver(
                 firstAtom,
                 firstCandidate,
                 forOverloadByLambdaReturnType = true,
+                withPCLASession = false,
             )
             while (iterator.hasNext()) {
                 val (candidate, atom) = iterator.next()

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

@@ -30,7 +30,6 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraint
 import org.jetbrains.kotlin.resolve.calls.model.PostponedAtomWithRevisableExpectedType
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
-import org.jetbrains.kotlin.types.model.TypeVariableTypeConstructorMarker
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.filterIsInstanceWithChecker
 
@@ -47,7 +46,7 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
         candidateReturnType: ConeKotlinType,
         context: ResolutionContext,
         collectVariablesFromContext: Boolean = false,
-        analyze: (PostponedResolvedAtom) -> Unit,
+        analyze: (PostponedResolvedAtom, Boolean) -> Unit,
     ) = c.runCompletion(completionMode, topLevelAtoms, candidateReturnType, context, collectVariablesFromContext, analyze)
 
     private fun ConstraintSystemCompletionContext.runCompletion(
@@ -56,7 +55,7 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
         topLevelType: ConeKotlinType,
         context: ResolutionContext,
         collectVariablesFromContext: Boolean = false,
-        analyze: (PostponedResolvedAtom) -> Unit,
+        analyze: (PostponedResolvedAtom, Boolean) -> Unit,
     ) {
         val topLevelTypeVariables = topLevelType.extractTypeVariables()
 
@@ -75,8 +74,10 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
             ) return
 
             // Stage 1: analyze postponed arguments with fixed parameter types
-            if (analyzeArgumentWithFixedParameterTypes(languageVersionSettings, postponedArguments, analyze))
+            if (analyzeArgumentWithFixedParameterTypes(languageVersionSettings, postponedArguments) {
-                continue
+                    analyze(it, /* withPCLASession = */ false)
+                }
+            ) continue
 
             val isThereAnyReadyForFixationVariable = findFirstVariableForFixation(
                 collectVariablesFromContext,
@@ -141,8 +142,10 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
             }
 
             // Stage 5: analyze the next ready postponed argument
-            if (analyzeNextReadyPostponedArgument(languageVersionSettings, postponedArguments, completionMode, analyze))
+            if (analyzeNextReadyPostponedArgument(languageVersionSettings, postponedArguments, completionMode) {
-                continue
+                    analyze(it, /* withPCLASession = */ false)
+                }
+            ) continue
 
             // Stage 6: fix next ready type variable with proper constraints
             if (fixNextReadyVariable(completionMode, topLevelAtoms, topLevelType, collectVariablesFromContext, postponedArguments))
@@ -159,8 +162,10 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
             if (completionMode == ConstraintSystemCompletionMode.PCLA_POSTPONED_CALL) {
                 // Complete all lambdas, maybe with fixing type variables used as top-level input types.
                 // It's necessary because we need to process all data-flow info before going to the next statement.
-                if (analyzeRemainingNotAnalyzedPostponedArgument(postponedArguments, analyze))
+                if (analyzeRemainingNotAnalyzedPostponedArgument(postponedArguments) {
-                    continue
+                        analyze(it, /* withPCLASession = */ false)
+                    }
+                ) continue
             }
 
             // Stage 8: report "not enough information" for uninferred type variables
@@ -176,8 +181,10 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
 
             // Stage 9: force analysis of remaining not analyzed postponed arguments and rerun stages if there are
             if (completionMode.allLambdasShouldBeAnalyzed) {
-                if (analyzeRemainingNotAnalyzedPostponedArgument(postponedArguments, analyze))
+                if (analyzeRemainingNotAnalyzedPostponedArgument(postponedArguments) {
-                    continue
+                        analyze(it, /* withPCLASession = */ false)
+                    }
+                ) continue
             }
 
             break
@@ -236,7 +243,7 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
     private fun ConstraintSystemCompletionContext.tryToCompleteWithPCLA(
         completionMode: ConstraintSystemCompletionMode,
         postponedArguments: List<PostponedResolvedAtom>,
-        analyze: (PostponedResolvedAtom) -> Unit,
+        analyze: (PostponedResolvedAtom, Boolean) -> Unit,
     ): Boolean {
         if (!completionMode.allLambdasShouldBeAnalyzed) return false
 
@@ -250,7 +257,7 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
             val notFixedInputTypeVariables = argument.inputTypes.flatMap { it.extractTypeVariables() }.filter { it !in fixedTypeVariables }
 
             if (notFixedInputTypeVariables.isEmpty()) continue
-            analyze(argument)
+            analyze(argument, /* withPCLASession = */ true)
 
             anyAnalyzed = true
         }

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

@@ -182,8 +182,8 @@ class FirCallCompleter(
             listOf(call),
             initialType,
             transformer.resolutionContext
-        ) {
+        ) { atom, withPCLASession ->
-            analyzer.analyze(candidate.system, it, candidate)
+            analyzer.analyze(candidate.system, atom, candidate, withPCLASession)
         }
     }
 

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

@@ -162,7 +162,7 @@ class FirDelegatedPropertyInferenceSession(
                 ConstraintSystemCompletionMode.FULL,
                 notCompletedCalls as List<FirStatement>,
                 unitType, resolutionContext
-            ) { lambdaAtom ->
+            ) { lambdaAtom, withPCLASession ->
                 // Reversed here bc we want top-most call to avoid exponential visit
                 val containingCandidateForLambda = notCompletedCalls.asReversed().first {
                     var found = false
@@ -177,6 +177,7 @@ class FirDelegatedPropertyInferenceSession(
                     parentSystem,
                     lambdaAtom,
                     containingCandidateForLambda,
+                    withPCLASession
                 )
             }
         }

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

@@ -55,16 +55,17 @@ class PostponedArgumentsAnalyzer(
         c: PostponedArgumentsAnalyzerContext,
         argument: PostponedResolvedAtom,
         candidate: Candidate,
+        withPCLASession: Boolean,
     ) {
         when (argument) {
             is ResolvedLambdaAtom ->
-                analyzeLambda(c, argument, candidate, forOverloadByLambdaReturnType = false)
+                analyzeLambda(c, argument, candidate, forOverloadByLambdaReturnType = false, withPCLASession)
 
             is LambdaWithTypeVariableAsExpectedTypeAtom ->
                 analyzeLambda(
                     c,
                     argument.transformToResolvedLambda(c.getBuilder(), resolutionContext),
-                    candidate, forOverloadByLambdaReturnType = false
+                    candidate, forOverloadByLambdaReturnType = false, withPCLASession
                 )
 
             is ResolvedCallableReferenceAtom -> processCallableReference(argument, candidate)
@@ -106,6 +107,7 @@ class PostponedArgumentsAnalyzer(
         lambda: ResolvedLambdaAtom,
         candidate: Candidate,
         forOverloadByLambdaReturnType: Boolean,
+        withPCLASession: Boolean,
         //diagnosticHolder: KotlinDiagnosticsHolder
     ): ReturnArgumentsAnalysisResult {
         // TODO: replace with `require(!lambda.analyzed)` when KT-54767 will be fixed
@@ -142,12 +144,6 @@ class PostponedArgumentsAnalyzer(
             else -> null
         }
 
-        val withPCLASession =
-            lambda.inputTypes
-                .any { inputType ->
-                    with(c) { inputType.extractTypeVariables() }.any(c.notFixedTypeVariables::contains)
-                }
-
         val results = lambdaAnalyzer.analyzeAndGetLambdaReturnArguments(
             lambda,
             receiver,