[FIR] Add contracts resolve for local classes

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.AdapterForResolvePhase
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.TransformImplicitType
+import org.jetbrains.kotlin.fir.resolve.transformers.contracts.runContractResolveForLocalClass
 import org.jetbrains.kotlin.fir.symbols.impl.FirAccessorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
@@ -52,13 +53,13 @@ class FirImplicitTypeBodyResolveTransformerAdapter(private val scopeSession: Sco
     }
 }
 
-fun <F : FirClass<F>> F.runBodiesResolutionForLocalClass(
+fun <F : FirClass<F>> F.runContractAndBodiesResolutionForLocalClass(
     components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
     resolutionMode: ResolutionMode,
     localClassesNavigationInfo: LocalClassesNavigationInfo,
 ): F {
     val (designationMap, targetedClasses) = localClassesNavigationInfo.run {
-        designationMap to parentForClass.keys + this@runBodiesResolutionForLocalClass
+        designationMap to parentForClass.keys + this@runContractAndBodiesResolutionForLocalClass
     }
 
     val implicitBodyResolveComputationSession =
@@ -74,6 +75,8 @@ fun <F : FirClass<F>> F.runBodiesResolutionForLocalClass(
     val newContext = components.context.createSnapshotForLocalClasses(returnTypeCalculator, targetedClasses)
     returnTypeCalculator.outerBodyResolveContext = newContext
 
+    runContractResolveForLocalClass(components.session, components.scopeSession, components.context, targetedClasses)
+
     val transformer = FirImplicitAwareBodyResolveTransformer(
         components.session, components.scopeSession,
         implicitBodyResolveComputationSession,

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.transformers.contracts.runContractResolveForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runStatusResolveForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runSupertypeResolvePhaseForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runTypeResolvePhaseForLocalClass
@@ -33,6 +34,6 @@ fun <F : FirClass<F>> F.runAllPhasesForLocalClass(
         components.createCurrentScopeList()
     )
     runStatusResolveForLocalClass(components.session)
-    runBodiesResolutionForLocalClass(components, resolutionMode, localClassesNavigationInfo)
+    runContractAndBodiesResolutionForLocalClass(components, resolutionMode, localClassesNavigationInfo)
     return this
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformer.kt

@@ -25,11 +25,16 @@ import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.compose
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 
-class FirContractResolveTransformer(session: FirSession, scopeSession: ScopeSession) : FirBodyResolveTransformer(
+class FirContractResolveTransformer(
+    session: FirSession,
+    scopeSession: ScopeSession,
+    outerBodyResolveContext: BodyResolveContext? = null
+) : FirBodyResolveTransformer(
     session,
     FirResolvePhase.CONTRACTS,
     implicitTypeOnly = false,
-    scopeSession
+    scopeSession,
+    outerBodyResolveContext = outerBodyResolveContext
 ) {
     override val declarationsTransformer: FirDeclarationsResolveTransformer = FirDeclarationsContractResolveTransformer(this)
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformerAdapter.kt

@@ -6,14 +6,20 @@
 package org.jetbrains.kotlin.fir.resolve.transformers.contracts
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.AdapterForResolvePhase
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.fir.visitors.transformSingle
 
 @AdapterForResolvePhase
 class FirContractResolveTransformerAdapter(private val scopeSession: ScopeSession) : FirTransformer<Nothing?>() {
@@ -29,3 +35,18 @@ class FirContractResolveTransformerAdapter(private val scopeSession: ScopeSessio
         return file.transform(transformer, ResolutionMode.ContextIndependent)
     }
 }
+
+fun <F : FirClass<F>> F.runContractResolveForLocalClass(
+    session: FirSession,
+    scopeSession: ScopeSession,
+    outerBodyResolveContext: FirAbstractBodyResolveTransformer.BodyResolveContext,
+    targetedClasses: Set<FirClass<*>>
+): F {
+    val newContext = outerBodyResolveContext.createSnapshotForLocalClasses(
+        ReturnTypeCalculatorForFullBodyResolve(),
+        targetedClasses
+    )
+    val transformer = FirContractResolveTransformer(session, scopeSession, newContext)
+
+    return this.transformSingle(transformer, ResolutionMode.ContextIndependent)
+}