FIR: Provide Candidate receiver to argument resolution helpers

The receiver might be useful for working with postponed atoms
and other info from the Candidate

^KT-32725 In Progress

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

@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 
 
-fun resolveArgumentExpression(
+fun Candidate.resolveArgumentExpression(
     /*
     csBuilder: ConstraintSystemBuilder,
     argument: KotlinCallArgument,
@@ -35,7 +35,6 @@ fun resolveArgumentExpression(
     sink: CheckerSink,
     isReceiver: Boolean,
     isSafeCall: Boolean,
-    acceptLambdaAtoms: (PostponedResolvedAtomMarker) -> Unit,
     typeProvider: (FirExpression) -> FirTypeRef?
 ) {
     return when (argument) {
@@ -61,7 +60,7 @@ fun resolveArgumentExpression(
             typeProvider
         )
         // TODO:!
-        is FirAnonymousFunction -> preprocessLambdaArgument(csBuilder, argument, expectedType, expectedTypeRef, acceptLambdaAtoms)
+        is FirAnonymousFunction -> preprocessLambdaArgument(csBuilder, argument, expectedType, expectedTypeRef)
         // TODO:!
         //TODO: Collection literal
         is FirWrappedArgumentExpression -> resolveArgumentExpression(
@@ -72,7 +71,6 @@ fun resolveArgumentExpression(
             sink,
             isReceiver,
             isSafeCall,
-            acceptLambdaAtoms,
             typeProvider
         )
         is FirBlock -> resolveBlockArgument(
@@ -83,14 +81,13 @@ fun resolveArgumentExpression(
             sink,
             isReceiver,
             isSafeCall,
-            acceptLambdaAtoms,
             typeProvider
         )
         else -> resolvePlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, isSafeCall, typeProvider)
     }
 }
 
-private fun resolveBlockArgument(
+private fun Candidate.resolveBlockArgument(
     csBuilder: ConstraintSystemBuilder,
     block: FirBlock,
     expectedType: ConeKotlinType,
@@ -98,7 +95,6 @@ private fun resolveBlockArgument(
     sink: CheckerSink,
     isReceiver: Boolean,
     isSafeCall: Boolean,
-    acceptLambdaAtoms: (PostponedResolvedAtomMarker) -> Unit,
     typeProvider: (FirExpression) -> FirTypeRef?
 ) {
     val returnArguments = block.returnExpressions()
@@ -123,7 +119,6 @@ private fun resolveBlockArgument(
             sink,
             isReceiver,
             isSafeCall,
-            acceptLambdaAtoms,
             typeProvider
         )
     }
@@ -230,7 +225,6 @@ internal fun Candidate.resolveArgument(
         sink,
         isReceiver,
         isSafeCall,
-        { this.postponedAtoms += it },
         typeProvider
     )
 }

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

@@ -18,12 +18,11 @@ import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
-fun preprocessLambdaArgument(
+fun Candidate.preprocessLambdaArgument(
     csBuilder: ConstraintSystemBuilder,
     argument: FirAnonymousFunction,
     expectedType: ConeKotlinType?,
     expectedTypeRef: FirTypeRef,
-    acceptLambdaAtoms: (PostponedResolvedAtomMarker) -> Unit,
     forceResolution: Boolean = false
 ) {
     if (expectedType != null && !forceResolution && csBuilder.isTypeVariable(expectedType)) {
@@ -42,7 +41,7 @@ fun preprocessLambdaArgument(
 //        csBuilder.addSubtypeConstraint(lambdaType, expectedType, ArgumentConstraintPosition(argument))
     }
 
-    acceptLambdaAtoms(resolvedArgument)
+    postponedAtoms += resolvedArgument
 }
 
 fun preprocessCallableReference(

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

@@ -33,9 +33,10 @@ interface LambdaAnalyzer {
 
 
 class PostponedArgumentsAnalyzer(
-    val lambdaAnalyzer: LambdaAnalyzer,
-    val typeProvider: (FirExpression) -> FirTypeRef?,
-    val components: InferenceComponents
+    private val lambdaAnalyzer: LambdaAnalyzer,
+    private val typeProvider: (FirExpression) -> FirTypeRef?,
+    private val components: InferenceComponents,
+    private val candidate: Candidate
 ) {
 
     fun analyze(
@@ -103,7 +104,7 @@ class PostponedArgumentsAnalyzer(
 
         val subResolvedKtPrimitives = returnArguments.map {
             var atom: PostponedResolvedAtomMarker? = null
-            resolveArgumentExpression(
+            candidate.resolveArgumentExpression(
                 c.getBuilder(),
                 it,
                 lambda.returnType.let(::substitute),
@@ -111,7 +112,6 @@ class PostponedArgumentsAnalyzer(
                 checkerSink,
                 false,
                 false,
-                { atom = it },
                 typeProvider
             )
 //            resolveKtPrimitive(

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

@@ -93,7 +93,7 @@ internal sealed class CheckReceivers : ResolutionStage() {
 
         if (expectedReceiverType != null) {
             if (explicitReceiverExpression != null && explicitReceiverKind.shouldBeCheckedAgainstExplicit()) {
-                resolveArgumentExpression(
+                candidate.resolveArgumentExpression(
                     candidate.csBuilder,
                     argument = explicitReceiverExpression,
                     expectedType = candidate.substitutor.substituteOrSelf(expectedReceiverType),
@@ -101,8 +101,7 @@ internal sealed class CheckReceivers : ResolutionStage() {
                     sink = sink,
                     isReceiver = true,
                     isSafeCall = callInfo.isSafeCall,
-                    typeProvider = callInfo.typeProvider,
-                    acceptLambdaAtoms = { candidate.postponedAtoms += it }
+                    typeProvider = callInfo.typeProvider
                 )
                 sink.yield()
             } else {

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

@@ -70,7 +70,8 @@ class FirCallCompleter(
         val completionMode = candidate.computeCompletionMode(inferenceComponents, expectedTypeRef, initialType)
         val replacements = mutableMapOf<FirExpression, FirExpression>()
 
-        val analyzer = PostponedArgumentsAnalyzer(LambdaAnalyzerImpl(replacements), { it.resultType }, inferenceComponents)
+        val analyzer =
+            PostponedArgumentsAnalyzer(LambdaAnalyzerImpl(replacements), { it.resultType }, inferenceComponents, candidate)
 
         completer.complete(candidate.system.asConstraintSystemCompleterContext(), completionMode, listOf(call), initialType) {
             analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it)
@@ -137,4 +138,4 @@ class FirCallCompleter(
             return (newLambdaExpression.body?.returnExpressions() ?: emptyList()) to InferenceSession.default
         }
     }
-}
+}