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
     }

compiler/testData/codegen/boxInline/callableReference/kt15449.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // FILE: 1.kt
 
 package test
@@ -38,4 +37,4 @@ class A  {
 fun box(): String {
     A().test()
     return result
-}
+}

compiler/testData/diagnostics/tests/callableReference/lambdaResult.fir.kt

@@ -1,43 +0,0 @@
-// SKIP_TXT
-// !DIAGNOSTICS: -UNUSED_PARAMETER
-// !LANGUAGE: +NewInference
-
-interface Inv<T>
-
-fun <E> Inv<E>.foo(
-    handler: () -> ((command: E) -> Unit)
-) {}
-
-fun bar(x: Int) {}
-fun bar(x: String) {}
-
-fun bar1(arg: Int) {}
-fun foo1(f: () -> (Int) -> Unit) = ""
-
-fun main(x: Inv<Int>) {
-    x.foo<Int> {
-        if (x.hashCode() == 0) return@foo <!UNRESOLVED_REFERENCE!>::bar<!>
-
-        ::bar
-    }
-
-    x.foo {
-        if (x.hashCode() == 0) return@foo <!UNRESOLVED_REFERENCE!>::bar<!>
-
-        ::bar
-    }
-
-    foo1 {
-        ::bar1
-    }
-
-    foo1 {
-        return@foo1 ::bar1
-    }
-
-    foo1 {
-        if (x.hashCode() == 0) return@foo1 ::bar
-
-        ::bar
-    }
-}

compiler/testData/diagnostics/tests/callableReference/lambdaResult.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // SKIP_TXT
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // !LANGUAGE: +NewInference