FIR: save/restore tower data contexts during resolve jumps

#KT-50278 Fixed

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerFirTestdataTestGenerated.java

@@ -5663,6 +5663,12 @@ public class DiagnosisCompilerFirTestdataTestGenerated extends AbstractDiagnosis
                 runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/receiverResolutionInLambda.kt");
             }
 
+            @Test
+            @TestMetadata("resolveToLocalFromCallSite.kt")
+            public void testResolveToLocalFromCallSite() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/resolveToLocalFromCallSite.kt");
+            }
+
             @Test
             @TestMetadata("selfReferenceToCompanionObject.kt")
             public void testSelfReferenceToCompanionObject() throws Exception {

compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/resolveToLocalFromCallSite.fir.txt

@@ -0,0 +1,24 @@
+FILE: resolveToLocalFromCallSite.kt
+    public final val x: R|kotlin/Any| = object : R|kotlin/Any| {
+        private constructor(): R|<anonymous>| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final fun foo(types: R|kotlin/collections/List<kotlin/String>|): R|kotlin/Unit| {
+            lval length: R|kotlin/String| = String(123)
+            R|<local>/types|.R|kotlin/collections/mapIndexed|<R|kotlin/String|, R|kotlin/Triple<@R|kotlin/ParameterName|(name = String(index))  kotlin/Int, kotlin/String, kotlin/Int>|>(<L> = mapIndexed@fun <anonymous>(i: R|@R|kotlin/ParameterName|(name = String(index))  kotlin/Int|, length: R|kotlin/String|): R|kotlin/Triple<@R|kotlin/ParameterName|(name = String(index))  kotlin/Int, kotlin/String, kotlin/Int>| <inline=Inline, kind=UNKNOWN>  {
+                ^ R|kotlin/Triple.Triple|<R|@R|kotlin/ParameterName|(name = String(index))  kotlin/Int|, R|kotlin/String|, R|kotlin/Int|>(R|<local>/i|, R|<local>/length|, (this@R|/<anonymous>|, R|<local>/length|).R|/<anonymous>.getFilteredType|())
+            }
+            )
+        }
+
+        private final fun R|kotlin/String|.getFilteredType(): R|kotlin/Int| {
+            ^getFilteredType R|/bar|(this@R|/<anonymous>.getFilteredType|.R|kotlin/String.length|)
+        }
+
+    }
+
+        public get(): R|kotlin/Any|
+    public final fun bar(x: R|kotlin/Int|): R|kotlin/Int| {
+        ^bar R|<local>/x|
+    }

compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/resolveToLocalFromCallSite.kt

@@ -0,0 +1,10 @@
+val x = object {
+    fun foo(types: List<String>) {
+        val length = "123"
+        types.mapIndexed { i, length -> Triple(i, length, length.getFilteredType()) }
+    }
+
+    private fun String.getFilteredType() = bar(length)
+}
+
+fun bar(x: Int) = x

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java

@@ -5663,6 +5663,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                 runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/receiverResolutionInLambda.kt");
             }
 
+            @Test
+            @TestMetadata("resolveToLocalFromCallSite.kt")
+            public void testResolveToLocalFromCallSite() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/resolveToLocalFromCallSite.kt");
+            }
+
             @Test
             @TestMetadata("selfReferenceToCompanionObject.kt")
             public void testSelfReferenceToCompanionObject() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -5663,6 +5663,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
                 runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/receiverResolutionInLambda.kt");
             }
 
+            @Test
+            @TestMetadata("resolveToLocalFromCallSite.kt")
+            public void testResolveToLocalFromCallSite() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/resolveToLocalFromCallSite.kt");
+            }
+
             @Test
             @TestMetadata("selfReferenceToCompanionObject.kt")
             public void testSelfReferenceToCompanionObject() throws Exception {

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.declarations.utils.hasExplicitBackingField
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
+import org.jetbrains.kotlin.fir.resolve.FirRegularTowerDataContexts
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
@@ -205,7 +206,14 @@ private class ReturnTypeCalculatorWithJump(
     ) -> FirDesignatedBodyResolveTransformerForReturnTypeCalculator = ::FirDesignatedBodyResolveTransformerForReturnTypeCalculator,
 ) : ReturnTypeCalculator {
 
+    @OptIn(PrivateForInline::class)
     var outerBodyResolveContext: BodyResolveContext? = null
+        set(context) {
+            field = context
+            outerTowerDataContexts = context?.regularTowerDataContexts
+        }
+
+    var outerTowerDataContexts: FirRegularTowerDataContexts? = null
 
     override fun tryCalculateReturnType(declaration: FirTypedDeclaration): FirResolvedTypeRef {
         if (declaration is FirValueParameter && declaration.returnTypeRef is FirImplicitTypeRef) {
@@ -253,6 +261,7 @@ private class ReturnTypeCalculatorWithJump(
         }
     }
 
+    @OptIn(PrivateForInline::class)
     private fun computeReturnTypeRef(declaration: FirCallableDeclaration): FirResolvedTypeRef {
         val symbol = declaration.symbol
         val provider = session.firProvider
@@ -277,6 +286,9 @@ private class ReturnTypeCalculatorWithJump(
             (listOf(file) + outerClasses.filterNotNull().asReversed()) to null
         }
 
+        val previousTowerDataContexts = outerBodyResolveContext?.regularTowerDataContexts
+        outerBodyResolveContext?.regularTowerDataContexts = outerTowerDataContexts!!
+
         val transformer = createTransformer(
             (designation.drop(1) + declaration).iterator(),
             session,
@@ -293,6 +305,9 @@ private class ReturnTypeCalculatorWithJump(
 
         val newReturnTypeRef = transformedDeclaration.returnTypeRef
         require(newReturnTypeRef is FirResolvedTypeRef) { transformedDeclaration.render() }
+        if (previousTowerDataContexts != null) {
+            outerBodyResolveContext.regularTowerDataContexts = previousTowerDataContexts
+        }
         return newReturnTypeRef
     }
 }