FIR: Fix callable references resolution when they're being returned from lambdas

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

@@ -263,13 +263,15 @@ class FirCallResolver(
             conflictResolver.chooseMaximallySpecificCandidates(bestCandidates, discriminateGenerics = false)
         }
 
+        resolvedCallableReferenceAtom.hasBeenResolvedOnce = true
+
         when {
             noSuccessfulCandidates -> {
                 return false
             }
             reducedCandidates.size > 1 -> {
-                if (resolvedCallableReferenceAtom.postponed) return false
-                resolvedCallableReferenceAtom.postponed = true
+                if (resolvedCallableReferenceAtom.hasBeenPostponed) return false
+                resolvedCallableReferenceAtom.hasBeenPostponed = true
                 return true
             }
         }

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

@@ -291,8 +291,9 @@ class ConstraintSystemCompleter(private val components: BodyResolveComponents) {
                             postponedAtom.collectNotFixedVariables()
                         }
                         postponedAtom is ResolvedCallableReferenceAtom -> {
-                            if (postponedAtom.postponed)
+                            if (postponedAtom.mightNeedAdditionalResolution) {
                                 postponedAtom.collectNotFixedVariables()
+                            }
                         }
                     }
                 }

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

@@ -68,7 +68,7 @@ class PostponedArgumentsAnalyzer(
     }
 
     private fun processCallableReference(atom: ResolvedCallableReferenceAtom, candidate: Candidate) {
-        if (atom.postponed) {
+        if (atom.mightNeedAdditionalResolution) {
             callResolver.resolveCallableReference(candidate.csBuilder, atom)
         }
 

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

@@ -116,35 +116,38 @@ class ResolvedCallableReferenceAtom(
     val lhs: DoubleColonLHS?,
     private val session: FirSession
 ) : PostponedResolvedAtom(), PostponedCallableReferenceMarker {
-    // TODO: in several places atoms are filtered by the marker interface - potential overhead/errors
-    var postponed: Boolean = false
+
+    var hasBeenResolvedOnce: Boolean = false
+    var hasBeenPostponed: Boolean = false
+
+    val mightNeedAdditionalResolution get() = !hasBeenResolvedOnce || hasBeenPostponed
 
     var resultingCandidate: Pair<Candidate, CandidateApplicability>? = null
 
     override val inputTypes: Collection<ConeKotlinType>
         get() {
-            if (!postponed) return emptyList()
+            if (!hasBeenPostponed) return emptyList()
             return extractInputOutputTypesFromCallableReferenceExpectedType(expectedType, session)?.inputTypes
                 ?: listOfNotNull(expectedType)
         }
     override val outputType: ConeKotlinType?
         get() {
-            if (!postponed) return null
+            if (!hasBeenPostponed) return null
             return extractInputOutputTypesFromCallableReferenceExpectedType(expectedType, session)?.outputType
         }
 
     override val expectedType: ConeKotlinType?
-        get() = if (!postponed)
+        get() = if (!hasBeenPostponed)
             initialExpectedType
         else
             revisedExpectedType ?: initialExpectedType
 
     override var revisedExpectedType: ConeKotlinType? = null
-        get() = if (postponed) field else expectedType
+        get() = if (hasBeenPostponed) field else expectedType
         private set
 
     override fun reviseExpectedType(expectedType: KotlinTypeMarker) {
-        if (!postponed) return
+        if (!mightNeedAdditionalResolution) return
         require(expectedType is ConeKotlinType)
         revisedExpectedType = expectedType
     }