[FIR] Get FirSession from declaration in ReturnTypeCalculatorWithJump

If the session is from the other module, it might not know about
FirFiles which are relevant for transforming and resolving the
declaration

^KT-52136 Fixed

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostics/fir/PersistentCheckerContextFactory.kt

@@ -15,7 +15,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.LLFirDesi
 internal object PersistentCheckerContextFactory {
     fun createEmptyPersistenceCheckerContext(sessionHolder: SessionHolder): PersistentCheckerContext {
         val returnTypeCalculator = createReturnTypeCalculatorForIDE(
-            sessionHolder.session,
             sessionHolder.scopeSession,
             ImplicitBodyResolveComputationSession(),
             ::LLFirDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirDesignatedBodyResolveTransformer.kt

@@ -36,7 +36,6 @@ internal class LLFirDesignatedBodyResolveTransformer(
     implicitTypeOnly = false,
     scopeSession = scopeSession,
     returnTypeCalculator = createReturnTypeCalculatorForIDE(
-        session,
         scopeSession,
         ImplicitBodyResolveComputationSession(),
         ::LLFirEnsureBasedTransformerForReturnTypeCalculator

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirDesignatedImplicitTypesTransformer.kt

@@ -38,7 +38,6 @@ internal class LLFirDesignatedImplicitTypesTransformer(
     scopeSession = scopeSession,
     firTowerDataContextCollector = towerDataContextCollector,
     returnTypeCalculator = createReturnTypeCalculatorForIDE(
-        session,
         scopeSession,
         implicitBodyResolveComputationSession,
         ::LLFirDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirFileAnnotationsResolveTransformer.kt

@@ -28,7 +28,6 @@ internal class LLFirFileAnnotationsResolveTransformer(
     implicitTypeOnly = false,
     scopeSession = scopeSession,
     returnTypeCalculator = createReturnTypeCalculatorForIDE(
-        session,
         scopeSession,
         implicitBodyResolveComputationSession,
         ::LLFirDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator

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

@@ -254,6 +254,12 @@ public class DiagnosisCompilerFirTestdataTestGenerated extends AbstractDiagnosis
             runTest("compiler/fir/analysis-tests/testData/resolve/genericReceiverPropertyOverride.kt");
         }
 
+        @Test
+        @TestMetadata("implicitTypeFromOtherModule.kt")
+        public void testImplicitTypeFromOtherModule() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/implicitTypeFromOtherModule.kt");
+        }
+
         @Test
         @TestMetadata("implicitTypeInFakeOverride.kt")
         public void testImplicitTypeInFakeOverride() throws Exception {

compiler/fir/analysis-tests/legacy-fir-tests/tests-gen/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -219,6 +219,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
         runTest("compiler/fir/analysis-tests/testData/resolve/genericReceiverPropertyOverride.kt");
     }
 
+    @TestMetadata("implicitTypeFromOtherModule.kt")
+    public void testImplicitTypeFromOtherModule() throws Exception {
+        runTest("compiler/fir/analysis-tests/testData/resolve/implicitTypeFromOtherModule.kt");
+    }
+
     @TestMetadata("implicitTypeInFakeOverride.kt")
     public void testImplicitTypeInFakeOverride() throws Exception {
         runTest("compiler/fir/analysis-tests/testData/resolve/implicitTypeInFakeOverride.kt");

compiler/fir/analysis-tests/testData/resolve/implicitTypeFromOtherModule.fir.txt

@@ -0,0 +1,30 @@
+Module: lib
+FILE: module_lib_implicitTypeFromOtherModule.kt
+    package dependency
+
+    public final class Other : R|kotlin/Any| {
+        public constructor(): R|dependency/Other| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final class Lib : R|kotlin/Any| {
+        public constructor(): R|dependency/Lib| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final fun test(): R|dependency/Lib| {
+        ^test R|dependency/Lib.Lib|()
+    }
+Module: main
+FILE: module_main_implicitTypeFromOtherModule.kt
+    package main
+
+    public final fun usage(): R|kotlin/Unit| {
+        R|main/take|(R|dependency/test|())
+    }
+    public final fun take(a: R|dependency/Lib|): R|kotlin/Unit| {
+    }
+    public final fun take(a: R|dependency/Other|): R|kotlin/Unit| {
+    }

compiler/fir/analysis-tests/testData/resolve/implicitTypeFromOtherModule.kt

@@ -0,0 +1,27 @@
+// SKIP_JAVAC
+// This directive is needed to skip this test in LazyBodyIsNotTouchedTilContractsPhaseTestGenerated,
+//  because it fails to parse module structure of multimodule test
+
+// MODULE: lib
+package dependency
+
+class Other
+
+class Lib
+
+fun test() = Lib()
+
+// MODULE: main(lib)
+package main
+
+import dependency.test
+import dependency.Lib
+import dependency.Other
+
+fun usage() {
+    take(test())
+}
+
+fun take(a: Lib) {}
+
+fun take(a: Other) {}

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

@@ -254,6 +254,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
             runTest("compiler/fir/analysis-tests/testData/resolve/genericReceiverPropertyOverride.kt");
         }
 
+        @Test
+        @TestMetadata("implicitTypeFromOtherModule.kt")
+        public void testImplicitTypeFromOtherModule() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/implicitTypeFromOtherModule.kt");
+        }
+
         @Test
         @TestMetadata("implicitTypeInFakeOverride.kt")
         public void testImplicitTypeInFakeOverride() throws Exception {

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

@@ -254,6 +254,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/genericReceiverPropertyOverride.kt");
         }
 
+        @Test
+        @TestMetadata("implicitTypeFromOtherModule.kt")
+        public void testImplicitTypeFromOtherModule() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/implicitTypeFromOtherModule.kt");
+        }
+
         @Test
         @TestMetadata("implicitTypeInFakeOverride.kt")
         public void testImplicitTypeInFakeOverride() throws Exception {

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

@@ -40,7 +40,7 @@ class FirImplicitTypeBodyResolveProcessor(
 @AdapterForResolveProcessor
 class FirImplicitTypeBodyResolveTransformerAdapter(session: FirSession, scopeSession: ScopeSession) : FirTransformer<Any?>() {
     private val implicitBodyResolveComputationSession = ImplicitBodyResolveComputationSession()
-    private val returnTypeCalculator = ReturnTypeCalculatorWithJump(session, scopeSession, implicitBodyResolveComputationSession)
+    private val returnTypeCalculator = ReturnTypeCalculatorWithJump(scopeSession, implicitBodyResolveComputationSession)
 
     private val transformer = FirImplicitAwareBodyResolveTransformer(
         session,
@@ -77,7 +77,6 @@ fun <F : FirClassLikeDeclaration> F.runContractAndBodiesResolutionForLocalClass(
         ((components.returnTypeCalculator as? ReturnTypeCalculatorWithJump)?.implicitBodyResolveComputationSession
             ?: ImplicitBodyResolveComputationSession())
     val returnTypeCalculator = ReturnTypeCalculatorWithJump(
-        components.session,
         components.scopeSession,
         implicitBodyResolveComputationSession,
         designationMap,
@@ -114,7 +113,6 @@ private fun ReturnTypeCalculator.getTransformerCreator() = when (this) {
 }
 
 fun createReturnTypeCalculatorForIDE(
-    session: FirSession,
     scopeSession: ScopeSession,
     implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
     createTransformer: (
@@ -126,7 +124,7 @@ fun createReturnTypeCalculatorForIDE(
         BodyResolveContext?
     ) -> FirDesignatedBodyResolveTransformerForReturnTypeCalculator
 ): ReturnTypeCalculator =
-    ReturnTypeCalculatorWithJump(session, scopeSession, implicitBodyResolveComputationSession, createTransformer = createTransformer)
+    ReturnTypeCalculatorWithJump(scopeSession, implicitBodyResolveComputationSession, createTransformer = createTransformer)
 
 open class FirImplicitAwareBodyResolveTransformer(
     session: FirSession,
@@ -193,7 +191,6 @@ open class FirImplicitAwareBodyResolveTransformer(
 }
 
 private class ReturnTypeCalculatorWithJump(
-    private val session: FirSession,
     private val scopeSession: ScopeSession,
     val implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
     val designationMapForLocalClasses: Map<FirCallableDeclaration, List<FirClassLikeDeclaration>> = mapOf(),
@@ -262,12 +259,14 @@ private class ReturnTypeCalculatorWithJump(
 
     @OptIn(PrivateForInline::class)
     private fun computeReturnTypeRef(declaration: FirCallableDeclaration): FirResolvedTypeRef {
+        // To properly transform and resolve declaration's type, we need to use its module's session
+        val session = declaration.moduleData.session
         val symbol = declaration.symbol
-        val provider = session.firProvider
 
         val (designation, outerBodyResolveContext) = if (declaration in designationMapForLocalClasses) {
             designationMapForLocalClasses.getValue(declaration) to outerBodyResolveContext
         } else {
+            val provider = session.firProvider
             val file = provider.getFirCallableContainerFile(symbol)
 
             val outerClasses = generateSequence(symbol.containingClass()?.classId) { classId ->