FIR: Use proper tower data contexts for postponed callable references

E.g., it should contain local scopes of containing lambdas

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

@@ -249,12 +249,18 @@ class FirCallResolver(
         )
         // No reset here!
         val localCollector = CandidateCollector(components, components.resolutionStageRunner)
-        val result = towerResolver.runResolver(
+
-            info,
+        val towerDataContext =
-            transformer.resolutionContext,
+            transformer.context.towerDataContextForCallableReferences[callableReferenceAccess] ?: transformer.context.towerDataContext
-            collector = localCollector,
+
-            manager = TowerResolveManager(localCollector),
+        val result = transformer.context.withTowerDataContext(towerDataContext) {
-        )
+            towerResolver.runResolver(
+                info,
+                transformer.resolutionContext,
+                collector = localCollector,
+                manager = TowerResolveManager(localCollector),
+            )
+        }
         val bestCandidates = result.bestCandidates()
         val noSuccessfulCandidates = !result.currentApplicability.isSuccess
         val reducedCandidates = if (noSuccessfulCandidates) {

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

@@ -74,6 +74,9 @@ class PostponedArgumentsAnalyzer(
 
         val callableReferenceAccess = atom.reference
         atom.analyzed = true
+
+        resolutionContext.bodyResolveContext.towerDataContextForCallableReferences.remove(callableReferenceAccess)
+
         val (resultingCandidate, applicability) = atom.resultingCandidate
             ?: Pair(null, CandidateApplicability.INAPPLICABLE)
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -9,6 +9,7 @@ import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContextsForClassParts
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataElement
@@ -53,6 +54,7 @@ class BodyResolveContext(
     var containers: PersistentList<FirDeclaration> = persistentListOf()
 
     val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
+    val towerDataContextForCallableReferences: MutableMap<FirCallableReferenceAccess, FirTowerDataContext> = mutableMapOf()
 
     @set:PrivateForInline
     var inferenceSession: FirInferenceSession = FirInferenceSession.DEFAULT
@@ -200,6 +202,7 @@ class BodyResolveContext(
     ): BodyResolveContext = BodyResolveContext(returnTypeCalculator, dataFlowAnalyzerContext, targetedLocalClasses, outerLocalClassForNested).apply {
         file = this@BodyResolveContext.file
         towerDataContextForAnonymousFunctions.putAll(this@BodyResolveContext.towerDataContextForAnonymousFunctions)
+        towerDataContextForCallableReferences.putAll(this@BodyResolveContext.towerDataContextForCallableReferences)
         containers = this@BodyResolveContext.containers
         towerDataContext = this@BodyResolveContext.towerDataContext
         anonymousFunctionsAnalyzedInDependentContext.addAll(this@BodyResolveContext.anonymousFunctionsAnalyzedInDependentContext)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt

@@ -51,6 +51,7 @@ open class FirBodyResolveTransformer(
     override fun transformFile(file: FirFile, data: ResolutionMode): CompositeTransformResult<FirFile> {
         checkSessionConsistency(file)
         context.cleanContextForAnonymousFunction()
+        context.towerDataContextForCallableReferences.clear()
         context.cleanDataFlowContext()
         @OptIn(PrivateForInline::class)
         context.file = file

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -664,6 +664,8 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             return resolvedReference.compose()
         }
 
+        context.towerDataContextForCallableReferences[callableReferenceAccess] = context.towerDataContext
+
         dataFlowAnalyzer.exitCallableReference(callableReferenceAccessWithTransformedLHS)
         return callableReferenceAccessWithTransformedLHS.compose()
     }

compiler/testData/codegen/box/javaInterop/notNullAssertions/localEntities.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // SKIP_JDK6
 // TARGET_BACKEND: JVM
 // WITH_RUNTIME