FIR: Support rewritten implicit types calculator in IDE

Unlike the compiler, in the IDE different files may have different
resolve state.

So using ReturnTypeCalculatorForFullBodyResolve may be not correct there

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

@@ -127,8 +127,9 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         override val integerLiteralTypeApproximator: IntegerLiteralTypeApproximationTransformer =
             IntegerLiteralTypeApproximationTransformer(symbolProvider, inferenceComponents.ctx)
         override val integerOperatorsTypeUpdater: IntegerOperatorsTypeUpdater = IntegerOperatorsTypeUpdater(integerLiteralTypeApproximator)
+
+        @PublishedApi
         internal var containerIfAny: FirDeclaration? = null
-            private set
 
         override var container: FirDeclaration
             get() = containerIfAny!!
@@ -136,7 +137,7 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
                 containerIfAny = value
             }
 
-        internal inline fun <T> withContainer(declaration: FirDeclaration, crossinline f: () -> T): T {
+        inline fun <T> withContainer(declaration: FirDeclaration, crossinline f: () -> T): T {
             val prevContainer = containerIfAny
             containerIfAny = declaration
             val result = f()

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

@@ -6,8 +6,6 @@
 package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
@@ -16,33 +14,6 @@ import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.compose
 
-class FirDesignatedBodyResolveTransformer(
-    private val designation: Iterator<FirElement>,
-    session: FirSession,
-    scopeSession: ScopeSession = ScopeSession(),
-    implicitTypeOnly: Boolean = true
-) : FirBodyResolveTransformer(
-    session,
-    phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
-    implicitTypeOnly = implicitTypeOnly,
-    scopeSession = scopeSession
-) {
-    override fun <E : FirElement> transformElement(element: E, data: ResolutionMode): CompositeTransformResult<E> {
-        if (designation.hasNext()) {
-            designation.next().visitNoTransform(this, data)
-            return element.compose()
-        }
-        return super.transformElement(element, data)
-    }
-
-    override fun transformDeclaration(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
-        return components.withContainer(declaration) {
-            declaration.replaceResolvePhase(transformerPhase)
-            transformElement(declaration, data)
-        }
-    }
-}
-
 @Deprecated("It is temp", level = DeprecationLevel.WARNING, replaceWith = ReplaceWith("TODO(\"что-то нормальное\")"))
 class FirBodyResolveTransformerAdapter : FirTransformer<Nothing?>() {
     private val scopeSession = ScopeSession()

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

@@ -49,6 +49,9 @@ class FirImplicitTypeBodyResolveTransformerAdapter : FirTransformer<Nothing?>()
     }
 }
 
+fun createReturnTypeCalculatorForIDE(session: FirSession, scopeSession: ScopeSession): ReturnTypeCalculator =
+    ReturnTypeCalculatorWithJump(session, scopeSession, ImplicitBodyResolveComputationSession())
+
 // TODO: This class is unused because it's effectively unneeded while we have mutable trees
 // Once we decide to switch to persistent trees this, it should be applied to each file
 // If the decision is to stay with mutable trees, this class should be removed

idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirResolutionApi.kt

@@ -8,21 +8,26 @@ package org.jetbrains.kotlin.idea.fir
 import com.intellij.openapi.progress.ProgressIndicatorProvider
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.FirProvider
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.getClassDeclaredCallableSymbols
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDesignatedBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.fir.resolve.transformers.runResolve
 import org.jetbrains.kotlin.fir.scopes.impl.FirPackageMemberScope
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
 import org.jetbrains.kotlin.fir.types.FirUserTypeRef
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.fir.visitors.compose
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -188,13 +193,41 @@ private fun FirDeclaration.runResolve(
     if (designation.all { it.resolvePhase >= toPhase }) {
         return
     }
-    val transformer = FirDesignatedBodyResolveTransformer(
+    val transformer = FirDesignatedBodyResolveTransformerForIDE(
         designation.iterator(), state.getSession(psi as KtElement),
         implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE
     )
     file.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
 }
 
+private class FirDesignatedBodyResolveTransformerForIDE(
+    private val designation: Iterator<FirElement>,
+    session: FirSession,
+    scopeSession: ScopeSession = ScopeSession(),
+    implicitTypeOnly: Boolean = true
+) : FirBodyResolveTransformer(
+    session,
+    phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+    implicitTypeOnly = implicitTypeOnly,
+    scopeSession = scopeSession,
+    returnTypeCalculator = createReturnTypeCalculatorForIDE(session, scopeSession)
+) {
+    override fun <E : FirElement> transformElement(element: E, data: ResolutionMode): CompositeTransformResult<E> {
+        if (designation.hasNext()) {
+            designation.next().visitNoTransform(this, data)
+            return element.compose()
+        }
+        return super.transformElement(element, data)
+    }
+
+    override fun transformDeclaration(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
+        return components.withContainer(declaration) {
+            declaration.replaceResolvePhase(transformerPhase)
+            transformElement(declaration, data)
+        }
+    }
+}
+
 fun KtElement.getOrBuildFir(
     state: FirModuleResolveState,
     phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE