FIR: extract BodyResolveContext.withFile

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
 import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
@@ -16,7 +17,9 @@ import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
 import org.jetbrains.kotlin.fir.resolve.dfa.PersistentFlow
 import org.jetbrains.kotlin.fir.resolve.inference.FirInferenceSession
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+import org.jetbrains.kotlin.fir.resolve.transformers.withScopeCleanup
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.createImportingScopes
 import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
@@ -29,11 +32,14 @@ class BodyResolveContext(
     val targetedLocalClasses: Set<FirClass<*>> = emptySet(),
     val outerLocalClassForNested: MutableMap<FirClassLikeSymbol<*>, FirClassLikeSymbol<*>> = mutableMapOf()
 ) {
-    val fileImportsScope: MutableList<FirScope> = mutableListOf()
+    private val mutableFileImportsScope: MutableList<FirScope> = mutableListOf()
+
+    val fileImportsScope: List<FirScope>
+        get() = mutableFileImportsScope
 
     @set:PrivateForInline
     lateinit var file: FirFile
-        internal set
+        private set
 
     @set:PrivateForInline
     var towerDataContextsForClassParts: FirTowerDataContextsForClassParts =
@@ -228,4 +234,25 @@ class BodyResolveContext(
             replaceTowerDataContext(this@BodyResolveContext.towerDataContext)
             anonymousFunctionsAnalyzedInDependentContext.addAll(this@BodyResolveContext.anonymousFunctionsAnalyzedInDependentContext)
         }
+
+    // WITH FirElement functions
+
+    internal inline fun <T> withFile(
+        file: FirFile,
+        session: FirSession,
+        scopeSession: ScopeSession,
+        crossinline f: () -> T
+    ): T {
+        clear()
+        @OptIn(PrivateForInline::class)
+        this.file = file
+        return withScopeCleanup(mutableFileImportsScope) {
+            withTowerDataCleanup {
+                val importingScopes = createImportingScopes(file, session, scopeSession)
+                mutableFileImportsScope += importingScopes
+                addNonLocalTowerDataElements(importingScopes.map { it.asTowerDataElement(isLocal = false) })
+                f()
+            }
+        }
+    }
 }

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

@@ -5,19 +5,19 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
-import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirTargetElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.asTowerDataElement
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
-import org.jetbrains.kotlin.fir.resolve.transformers.withScopeCleanup
 import org.jetbrains.kotlin.fir.scopes.FirCompositeScope
-import org.jetbrains.kotlin.fir.scopes.createImportingScopes
 import org.jetbrains.kotlin.fir.scopes.impl.createCurrentScopeList
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.fir.types.FirTypeRef
 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.name.FqName
 
 open class FirBodyResolveTransformer(
     session: FirSession,
@@ -49,19 +48,10 @@ open class FirBodyResolveTransformer(
 
     override fun transformFile(file: FirFile, data: ResolutionMode): CompositeTransformResult<FirFile> {
         checkSessionConsistency(file)
-        context.clear()
-        @OptIn(PrivateForInline::class)
-        context.file = file
-        return withScopeCleanup(context.fileImportsScope) {
-            context.withTowerDataCleanup {
-                val importingScopes = createImportingScopes(file, session, components.scopeSession)
-                context.fileImportsScope += importingScopes
-                context.addNonLocalTowerDataElements(importingScopes.map { it.asTowerDataElement(isLocal = false) })
-
-                file.replaceResolvePhase(transformerPhase)
-                @Suppress("UNCHECKED_CAST")
-                transformDeclarationContent(file, data) as CompositeTransformResult<FirFile>
-            }
+        return context.withFile(file, session, components.scopeSession) {
+            file.replaceResolvePhase(transformerPhase)
+            @Suppress("UNCHECKED_CAST")
+            transformDeclarationContent(file, data) as CompositeTransformResult<FirFile>
         }
     }