[FIR IDE] Refactoring of designated IDE transformers

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirDesignationState.kt

@@ -8,10 +8,11 @@ package org.jetbrains.kotlin.idea.fir.low.level.api
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 
-abstract class ContextByDesignationCollector<C : Any>(private val designation: FirDeclarationDesignation, var firFile: FirFile) {
+abstract class ContextByDesignationCollector<C : Any>(private val designation: FirDeclarationDesignation) {
     private var context: C? = null
-    private val designationState = FirDesignationState(designation, firFile)
+    private val designationState = FirDesignationState(designation)
 
     protected abstract fun getCurrentContext(): C
     protected abstract fun goToNestedDeclaration(declaration: FirDeclaration)
@@ -37,7 +38,7 @@ abstract class ContextByDesignationCollector<C : Any>(private val designation: F
     }
 }
 
-private class FirDesignationState(val designation: FirDeclarationDesignation, val firFile: FirFile) {
+private class FirDesignationState(val designation: FirDeclarationDesignation) {
     /**
      * Holds current declaration index
      * if `currentIndex in [0, designation.path.lastIndex]` then current declaration is in path

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/DeclarationCopyBuilder.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/FirDeclarationDesignation.kt

@@ -12,71 +12,86 @@ import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
 
-class FirDeclarationDesignation(
+class FirDeclarationDesignationWithFile(
+    path: List<FirDeclaration>,
+    declaration: FirDeclaration,
+    isLocalDesignation: Boolean,
+    val firFile: FirFile
+) : FirDeclarationDesignation(
+    path,
+    declaration,
+    isLocalDesignation
+) {
+    fun toSequenceWithFile(includeTarget: Boolean): Sequence<FirDeclaration> = sequence {
+        yield(firFile)
+        yieldAll(path)
+        if (includeTarget) yield(declaration)
+    }
+}
+
+open class FirDeclarationDesignation(
     val path: List<FirDeclaration>,
     val declaration: FirDeclaration,
-    val isLocalDesignation: Boolean,
+    val isLocalDesignation: Boolean
 ) {
-    @OptIn(ExperimentalStdlibApi::class)
-    val fullDesignation = buildList {
-        addAll(path)
-        add(declaration)
+    fun toSequence(includeTarget: Boolean): Sequence<FirDeclaration> = sequence {
+        yieldAll(path)
+        if (includeTarget) yield(declaration)
     }
 }
 
+private fun collectDesignationAndIsLocal(declaration: FirDeclaration): Pair<List<FirDeclaration>, Boolean> {
+    val firProvider = declaration.moduleData.session.firProvider
 
-object DeclarationDesignationCollector {
-    fun collectDesignation(declaration: FirDeclaration): FirDeclarationDesignation {
-        val firProvider = declaration.moduleData.session.firProvider
+    val containingClass = when (declaration) {
+        is FirCallableDeclaration<*> -> declaration.containingClass()?.toFirRegularClass(declaration.moduleData.session)
+        is FirClassLikeDeclaration<*> -> declaration.symbol.classId.outerClassId?.let(firProvider::getFirClassifierByFqName)
+        else -> error("Invalid declaration ${declaration.renderWithType()}")
+    } ?: return emptyList<FirDeclaration>() to false
 
-        val containingClass = when (declaration) {
-            is FirCallableDeclaration<*> -> declaration.containingClass()?.toFirRegularClass(declaration.moduleData.session)
-            is FirClassLikeDeclaration<*> -> declaration.symbol.classId.outerClassId?.let(firProvider::getFirClassifierByFqName)
-            else -> error("Invalid declaration ${declaration.renderWithType()}")
-        } ?: return FirDeclarationDesignation(emptyList(), declaration, isLocalDesignation = false)
-
-        require(containingClass is FirRegularClass) {
-            "FirRegularClass as containing declaration expected but found ${containingClass.renderWithType()}"
-        }
-
-        val path = when {
-            containingClass.isLocal -> containingClass.collectForLocal()
-            else -> containingClass.collectForNonLocal()
-        }
-
-        return FirDeclarationDesignation(
-            path.reversed(),
-            declaration,
-            isLocalDesignation = false
-        )
+    require(containingClass is FirRegularClass) {
+        "FirRegularClass as containing declaration expected but found ${containingClass.renderWithType()}"
     }
 
-    private fun FirRegularClass.collectForNonLocal(): List<FirClassLikeDeclaration<*>> {
-        require(!isLocal)
-        val firProvider = moduleData.session.firProvider
-        var containingClassId = classId.outerClassId
-        val designation = mutableListOf<FirClassLikeDeclaration<*>>(this)
-        while (containingClassId != null) {
-            val currentClass = firProvider.getFirClassifierByFqName(containingClassId) ?: break
-            designation.add(currentClass)
-            containingClassId = containingClassId.outerClassId
-        }
-        return designation
-    }
-
-    private fun FirRegularClass.collectForLocal(): List<FirClassLikeDeclaration<*>> {
-        require(isLocal)
-        var containingClassLookUp = containingClassForLocal()
-        val designation = mutableListOf<FirClassLikeDeclaration<*>>(this)
-        while (containingClassLookUp != null && containingClassLookUp.classId.isLocal) {
-            val currentClass = containingClassLookUp.toFirRegularClass(moduleData.session) ?: break
-            designation.add(currentClass)
-            containingClassLookUp = currentClass.containingClassForLocal()
-        }
-        return designation
+    val path = when {
+        containingClass.isLocal -> containingClass.collectForLocal()
+        else -> containingClass.collectForNonLocal()
     }
+    return path.reversed() to containingClass.isLocal
 }
 
+private fun FirRegularClass.collectForNonLocal(): List<FirDeclaration> {
+    require(!isLocal)
+    val firProvider = moduleData.session.firProvider
+    var containingClassId = classId.outerClassId
+    val designation = mutableListOf<FirDeclaration>(this)
+    while (containingClassId != null) {
+        val currentClass = firProvider.getFirClassifierByFqName(containingClassId) ?: break
+        designation.add(currentClass)
+        containingClassId = containingClassId.outerClassId
+    }
+    return designation
+}
+
+private fun FirRegularClass.collectForLocal(): List<FirClassLikeDeclaration<*>> {
+    require(isLocal)
+    var containingClassLookUp = containingClassForLocal()
+    val designation = mutableListOf<FirClassLikeDeclaration<*>>(this)
+    while (containingClassLookUp != null && containingClassLookUp.classId.isLocal) {
+        val currentClass = containingClassLookUp.toFirRegularClass(moduleData.session) ?: break
+        designation.add(currentClass)
+        containingClassLookUp = currentClass.containingClassForLocal()
+    }
+    return designation
+}
+
+
+fun FirDeclaration.collectDesignation(firFile: FirFile): FirDeclarationDesignationWithFile =
+    collectDesignationAndIsLocal(this).let {
+        FirDeclarationDesignationWithFile(it.first, this, it.second, firFile)
+    }
 
 fun FirDeclaration.collectDesignation(): FirDeclarationDesignation =
-    DeclarationDesignationCollector.collectDesignation(this)
+    collectDesignationAndIsLocal(this).let {
+        FirDeclarationDesignation(it.first, this, it.second)
+    }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/LowLevelFirApiFacadeForResolveOnAir.kt

@@ -61,16 +61,14 @@ object LowLevelFirApiFacadeForResolveOnAir {
         }
     }
 
-    private fun <T : KtElement> onAirResolveElement(
+    fun <T : KtElement> onAirResolveElement(
         state: FirModuleResolveState,
         place: T,
         elementToResolve: T,
-    ): FirModuleResolveState {
+    ): FirElement {
         require(state is FirModuleResolveStateImpl)
         require(place.isPhysical)
-        require(!elementToResolve.isPhysical)
 
-        val collector = FirTowerDataContextAllElementsCollector()
         val declaration = runResolveBodyResolveOnAir(
             state = state,
             replacement = RawFirReplacement(place, elementToResolve),
@@ -88,10 +86,9 @@ object LowLevelFirApiFacadeForResolveOnAir {
                 element.acceptChildren(this)
             }
         }
-        declaration.accept(expressionLocator)
 
-        val recordedMap = FirElementsRecorder.recordElementsFrom(declaration, FirElementsRecorder())
-        return FirModuleResolveStateDepended(state, collector, recordedMap)
+        declaration.accept(expressionLocator)
+        return expressionLocator.result ?: error("Resolved on-air element was not found in containing declaration")
     }
 
     fun onAirGetTowerContextProvider(

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/fir/ContextCollectingDiagnosticCollectorVisitor.kt

@@ -13,16 +13,16 @@ import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.SessionHolder
 import org.jetbrains.kotlin.idea.fir.low.level.api.ContextByDesignationCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
 
 private class ContextCollectingDiagnosticCollectorVisitor private constructor(
     sessionHolder: SessionHolder,
-    designation: FirDeclarationDesignation,
-    firFile: FirFile,
+    designation: FirDeclarationDesignationWithFile,
 ) : AbstractDiagnosticCollectorVisitor(
     PersistentCheckerContextFactory.createEmptyPersistenceCheckerContext(sessionHolder)
 ) {
-    private val contextCollector = object : ContextByDesignationCollector<PersistentCheckerContext>(designation, firFile) {
+    private val contextCollector = object : ContextByDesignationCollector<PersistentCheckerContext>(designation) {
         override fun getCurrentContext(): PersistentCheckerContext = context
 
         override fun goToNestedDeclaration(declaration: FirDeclaration) {
@@ -41,9 +41,9 @@ private class ContextCollectingDiagnosticCollectorVisitor private constructor(
     override fun checkElement(element: FirElement) {}
 
     companion object {
-        fun collect(sessionHolder: SessionHolder, firFile: FirFile, designation: FirDeclarationDesignation): PersistentCheckerContext {
-            val visitor = ContextCollectingDiagnosticCollectorVisitor(sessionHolder, designation, firFile)
-            firFile.accept(visitor, null)
+        fun collect(sessionHolder: SessionHolder, designation: FirDeclarationDesignationWithFile): PersistentCheckerContext {
+            val visitor = ContextCollectingDiagnosticCollectorVisitor(sessionHolder, designation)
+            designation.firFile.accept(visitor, null)
             return visitor.contextCollector.getCollectedContext()
         }
     }
@@ -63,10 +63,10 @@ internal object PersistenceContextCollector {
         require(!isLocal) {
             "Cannot collect context for local declaration ${declaration.renderWithType()}"
         }
-        val designation = declaration.collectDesignation()
+        val designation = declaration.collectDesignation(firFile)
         check(!designation.isLocalDesignation) {
             "Designation should not local for ${declaration.renderWithType()}"
         }
-        return ContextCollectingDiagnosticCollectorVisitor.collect(sessionHolder, firFile, designation)
+        return ContextCollectingDiagnosticCollectorVisitor.collect(sessionHolder, designation)
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -24,7 +24,7 @@ internal object FirLazyBodiesCalculator {
     fun calculateLazyBodiesInside(element: FirElement, designation: FirDeclarationDesignation) {
         element.transform<FirElement, MutableList<FirDeclaration>>(
             FirLazyBodiesCalculatorTransformer,
-            designation.fullDesignation.toMutableList()
+            designation.toSequence(includeTarget = true).toMutableList()
         )
     }
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
@@ -143,7 +144,7 @@ internal class FirLazyDeclarationResolver(
         resolveFileAnnotations(containerFirFile, moduleFileCache, scopeSession)
 
         val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
-        val designation = nonLocalDeclarationToResolve.collectDesignation()
+        val designation = nonLocalDeclarationToResolve.collectDesignation(containerFirFile)
 
         executeWithoutPCE {
             calculateLazyBodies(firDeclarationToResolve, designation)
@@ -156,7 +157,6 @@ internal class FirLazyDeclarationResolver(
             if (currentPhase.pluginPhase) continue
             if (checkPCE) checkCanceled()
             runLazyResolvePhase(
-                containerFirFile,
                 currentPhase,
                 scopeSession,
                 towerDataContextCollector,
@@ -167,20 +167,18 @@ internal class FirLazyDeclarationResolver(
     }
 
     private fun runLazyResolvePhase(
-        containerFirFile: FirFile,
         phase: FirResolvePhase,
         scopeSession: ScopeSession,
         towerDataContextCollector: FirTowerDataContextCollector?,
         firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
-        designation: FirDeclarationDesignation
+        designation: FirDeclarationDesignationWithFile
     ) {
-        if (designation.fullDesignation.all { it.resolvePhase >= phase }) {
+        if (designation.toSequence(includeTarget = true).all { it.resolvePhase >= phase }) {
             return
         }
 
         val transformer = phase.createLazyTransformer(
             designation,
-            containerFirFile,
             scopeSession,
             towerDataContextCollector,
             firProviderInterceptor,
@@ -192,46 +190,41 @@ internal class FirLazyDeclarationResolver(
     }
 
     private fun FirResolvePhase.createLazyTransformer(
-        designation: FirDeclarationDesignation,
-        containerFirFile: FirFile,
+        designation: FirDeclarationDesignationWithFile,
         scopeSession: ScopeSession,
         towerDataContextCollector: FirTowerDataContextCollector?,
         firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
     ): FirLazyTransformerForIDE = when (this) {
         FirResolvePhase.SUPER_TYPES -> FirDesignatedSupertypeResolverTransformerForIDE(
             designation,
-            containerFirFile.moduleData.session,
+            designation.firFile.moduleData.session,
             scopeSession,
             firProviderInterceptor,
         )
         FirResolvePhase.SEALED_CLASS_INHERITORS -> FirLazyTransformerForIDE.EMPTY
         FirResolvePhase.TYPES -> FirDesignatedTypeResolverTransformerForIDE(
-            containerFirFile,
             designation,
-            containerFirFile.moduleData.session,
+            designation.firFile.moduleData.session,
             scopeSession,
         )
         FirResolvePhase.STATUS -> FirDesignatedStatusResolveTransformerForIDE(
             designation,
-            containerFirFile.moduleData.session,
+            designation.firFile.moduleData.session,
             scopeSession
         )
         FirResolvePhase.CONTRACTS -> FirDesignatedContractsResolveTransformerForIDE(
-            containerFirFile,
             designation,
-            containerFirFile.moduleData.session,
+            designation.firFile.moduleData.session,
             scopeSession,
         )
         FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE -> FirDesignatedImplicitTypesTransformerForIDE(
-            containerFirFile,
             designation,
-            containerFirFile.moduleData.session,
+            designation.firFile.moduleData.session,
             scopeSession
         )
         FirResolvePhase.BODY_RESOLVE -> FirDesignatedBodyResolveTransformerForIDE(
-            containerFirFile,
             designation,
-            containerFirFile.moduleData.session,
+            designation.firFile.moduleData.session,
             scopeSession,
             towerDataContextCollector,
             firProviderInterceptor

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirDesignatedBodyResolveTransformerForIDE.kt

@@ -14,11 +14,11 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
 
 internal class FirDesignatedBodyResolveTransformerForIDE(
-    private val firFile: FirFile,
-    designation: FirDeclarationDesignation,
+    private val designation: FirDeclarationDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
     towerDataContextCollector: FirTowerDataContextCollector? = null,
@@ -37,19 +37,19 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
     firTowerDataContextCollector = towerDataContextCollector,
     firProviderInterceptor = firProviderInterceptor,
 ) {
-    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation.toDesignationIterator())
+    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
     @Suppress("NAME_SHADOWING")
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration =
-        ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) { declaration, data ->
+        ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) {
             super.transformDeclarationContent(declaration, data)
         }
 
     override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
-        ideDeclarationTransformer.needReplacePhase(firDeclaration)
+        ideDeclarationTransformer.needReplacePhase
 
     override fun transformDeclaration() {
-        firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextDependent)
+        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextDependent)
     }
 }
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirDesignatedContractsResolveTransformerForIDE.kt

@@ -11,26 +11,25 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveTransformer
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 
 internal class FirDesignatedContractsResolveTransformerForIDE(
-    private val firFile: FirFile,
-    designation: FirDeclarationDesignation,
+    private val designation: FirDeclarationDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
 ) : FirLazyTransformerForIDE, FirContractResolveTransformer(session, scopeSession) {
-    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation.toDesignationIterator())
+    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
     @Suppress("NAME_SHADOWING")
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration =
-        ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) { declaration, data ->
+        ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) {
             super.transformDeclarationContent(declaration, data)
         }
 
 
-    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = ideDeclarationTransformer.needReplacePhase(firDeclaration)
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = ideDeclarationTransformer.needReplacePhase
 
     override fun transformDeclaration() {
-        firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextDependent)
+        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextDependent)
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirDesignatedImplicitTypesTransformerForIDE.kt

@@ -14,12 +14,11 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirImplicitAwareBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
 
 internal class FirDesignatedImplicitTypesTransformerForIDE(
-    private val firFile: FirFile,
-    designation: FirDeclarationDesignation,
+    private val designation: FirDeclarationDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
     implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession = ImplicitBodyResolveComputationSession(),
@@ -36,17 +35,17 @@ internal class FirDesignatedImplicitTypesTransformerForIDE(
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
     )
 ) {
-    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation.toDesignationIterator())
+    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
     @Suppress("NAME_SHADOWING")
     override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): FirDeclaration =
-        ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) { declaration, data ->
+        ideDeclarationTransformer.transformDeclarationContent(this, declaration, data) {
             super.transformDeclarationContent(declaration, data)
         }
 
-    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = ideDeclarationTransformer.needReplacePhase(firDeclaration)
+    override fun needReplacePhase(firDeclaration: FirDeclaration): Boolean = ideDeclarationTransformer.needReplacePhase
 
     override fun transformDeclaration() {
-        firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextDependent)
+        designation.firFile.transform<FirFile, ResolutionMode>(this, ResolutionMode.ContextDependent)
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirDesignatedStatusResolveTransformerForIDE.kt

@@ -13,27 +13,43 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirDesignatedStatusResolveT
 import org.jetbrains.kotlin.fir.resolve.transformers.FirStatusResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.StatusComputationSession
 import org.jetbrains.kotlin.fir.visitors.transformSingle
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePhase
 
 class FirDesignatedStatusResolveTransformerForIDE(
-    private val designation: FirDeclarationDesignation,
+    private val designation: FirDeclarationDesignationWithFile,
     private val session: FirSession,
     private val scopeSession: ScopeSession,
 ): FirLazyTransformerForIDE {
 
+    private val firstItemInDesignation = designation.path.firstOrNull() ?: designation.declaration
+
     private fun resolveClass(targetClass: FirClass<*>) {
         val transformer = FirDesignatedStatusResolveTransformer(
             session = session,
             scopeSession = scopeSession,
-            designation = designation.fullDesignation.iterator(),
+            designation = designation.toSequence(includeTarget = true).iterator(),
             targetClass = targetClass,
             statusComputationSession = StatusComputationSession.Regular(),
             designationMapForLocalClasses = emptyMap(),
-            scopeForLocalClass = null)
+            scopeForLocalClass = null
+        )
 
-        val firstClass = designation.fullDesignation.first()
-        firstClass.transformSingle(transformer, null)
+        firstItemInDesignation.transformSingle(transformer, null)
+    }
+
+    private fun resolveTypeAlias(targetClass: FirTypeAlias) {
+        val transformer = FirDesignatedStatusResolveTransformer(
+            session = session,
+            scopeSession = scopeSession,
+            designation = designation.toSequence(includeTarget = true).iterator(),
+            targetClass = targetClass,
+            statusComputationSession = StatusComputationSession.Regular(),
+            designationMapForLocalClasses = emptyMap(),
+            scopeForLocalClass = null
+        )
+
+        firstItemInDesignation.transformSingle(transformer, null)
     }
 
     private fun resolveTopLevelMethod(targetCallable: FirCallableDeclaration<*>) {
@@ -50,11 +66,12 @@ class FirDesignatedStatusResolveTransformerForIDE(
         val transformer = object : FirDesignatedStatusResolveTransformer(
             session = session,
             scopeSession = scopeSession,
-            designation = designation.fullDesignation.iterator(),
+            designation = designation.toSequence(includeTarget = true).iterator(),
             targetClass = containingClass,
             statusComputationSession = StatusComputationSession.Regular(),
             designationMapForLocalClasses = emptyMap(),
-            scopeForLocalClass = null) {
+            scopeForLocalClass = null
+        ) {
 
             override fun <F : FirClass<F>> transformClass(
                 klass: FirClass<F>,
@@ -66,25 +83,25 @@ class FirDesignatedStatusResolveTransformerForIDE(
             }
         }
 
-        val firstClass = designation.fullDesignation.first()
-        firstClass.transformSingle(transformer, null)
+        firstItemInDesignation.transformSingle(transformer, null)
     }
 
     override fun transformDeclaration() {
-        designation.ensurePhase(FirResolvePhase.STATUS, exceptLast = true)
-        when(val resolveTarget = designation.declaration) {
+        designation.ensurePhase(FirResolvePhase.STATUS, exceptTarget = true)
+        when (val resolveTarget = designation.declaration) {
             is FirClass<*> -> resolveClass(resolveTarget)
+            is FirTypeAlias -> resolveTypeAlias(resolveTarget)
             is FirCallableDeclaration<*> -> {
                 val containingClass = designation.path.lastOrNull()
                 if (containingClass == null) {
-                    check(designation.fullDesignation.size == 1) { "Invalid designation - should be single element designation for top level declaration" }
+                    check(designation.path.isEmpty()) { "Invalid designation - should be single element designation for top level declaration" }
                     resolveTopLevelMethod(resolveTarget)
                 } else {
                     check(containingClass is FirClass<*>) { "Invalid designation - the parent of callable is not a class" }
                     resolveClassMember(containingClass, resolveTarget)
                 }
             }
-            else -> error("Declaration should be ${FirClass::class.simpleName} but given ${resolveTarget::class.simpleName}")
+            else -> error("Unexpected declaration to resolve ${resolveTarget::class.simpleName}")
         }
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirDesignatedSupertypeResolverTransformerForIDE.kt

@@ -14,10 +14,11 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSu
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSupertypeResolverVisitor
 import org.jetbrains.kotlin.fir.resolve.transformers.SupertypeComputationSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePhase
 
 internal class FirDesignatedSupertypeResolverTransformerForIDE(
-    private val designation: FirDeclarationDesignation,
+    private val designation: FirDeclarationDesignationWithFile,
     private val session: FirSession,
     private val scopeSession: ScopeSession,
     private val firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
@@ -26,7 +27,7 @@ internal class FirDesignatedSupertypeResolverTransformerForIDE(
     private val supertypeComputationSession = SupertypeComputationSession()
 
     override fun transformDeclaration() {
-        designation.ensurePhase(FirResolvePhase.SUPER_TYPES, exceptLast = true)
+        designation.ensurePhase(FirResolvePhase.SUPER_TYPES, exceptTarget = true)
 
         val resolver = FirSupertypeResolverVisitor(
             session = session,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirDesignatedTypeResolverTransformerForIDE.kt

@@ -10,34 +10,33 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTypeResolveTransformer
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE.Companion.ensurePhase
 
 class FirDesignatedTypeResolverTransformerForIDE(
-    private val originalFile: FirFile,
-    private val designation: FirDeclarationDesignation,
+    private val designation: FirDeclarationDesignationWithFile,
     session: FirSession,
     scopeSession: ScopeSession,
-): FirLazyTransformerForIDE, FirTypeResolveTransformer(session, scopeSession) {
+) : FirLazyTransformerForIDE, FirTypeResolveTransformer(session, scopeSession) {
 
-    private val ideDeclarationTransformer = IDEDeclarationTransformer(FirDesignationIterator(designation.fullDesignation))
+    private val ideDeclarationTransformer = IDEDeclarationTransformer(designation)
 
     @Suppress("NAME_SHADOWING")
     override fun transformRegularClass(regularClass: FirRegularClass, data: Any?): FirStatement {
-        return ideDeclarationTransformer.transformDeclarationContent(this, regularClass, data) { klass, data ->
-            super.transformRegularClass(klass, data) as FirRegularClass
+        return ideDeclarationTransformer.transformDeclarationContent(this, regularClass, data) {
+            super.transformRegularClass(regularClass, data) as FirRegularClass
         }
     }
 
     @Suppress("NAME_SHADOWING")
     override fun transformAnonymousObject(anonymousObject: FirAnonymousObject, data: Any?): FirStatement {
-        return ideDeclarationTransformer.transformDeclarationContent(this, anonymousObject, data) { anonymousObject, data ->
+        return ideDeclarationTransformer.transformDeclarationContent(this, anonymousObject, data) {
             super.transformAnonymousObject(anonymousObject, data) as FirAnonymousObject
         }
     }
 
     override fun transformDeclaration() {
-        designation.ensurePhase(FirResolvePhase.TYPES, exceptLast = true)
-        originalFile.transform<FirFile, Any?>(this, null)
+        designation.ensurePhase(FirResolvePhase.TYPES, exceptTarget = true)
+        designation.firFile.transform<FirFile, Any?>(this, null)
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirDesignationIterator.kt

@@ -1,30 +0,0 @@
-/*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
-
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
-
-internal fun FirDeclarationDesignation.toDesignationIterator() =
-    FirDesignationIterator(fullDesignation)
-
-internal class FirDesignationIterator(val designation: List<FirDeclaration>) {
-    private var currentIndex = 0
-
-    fun canGoNext(): Boolean = currentIndex <= designation.lastIndex
-
-    fun isTargetDeclaration(): Boolean = currentIndex == designation.lastIndex
-
-    val currentDeclaration get() = designation[currentIndex]
-
-    fun goNext() {
-        if (canGoNext()) {
-            currentIndex++
-        } else {
-            throw IndexOutOfBoundsException()
-        }
-    }
-}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/FirLazyTransformerForIDE.kt

@@ -6,7 +6,7 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
 
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignationWithFile
 
 internal interface FirLazyTransformerForIDE {
     fun transformDeclaration()
@@ -16,10 +16,10 @@ internal interface FirLazyTransformerForIDE {
             override fun transformDeclaration() = Unit
         }
 
-        fun FirDeclarationDesignation.ensurePhase(firResolvePhase: FirResolvePhase, exceptLast: Boolean) {
-            val designationList = if (exceptLast) path else fullDesignation
-            designationList.forEach { firDeclaration ->
-                check (firDeclaration.resolvePhase >= firResolvePhase) {
+        fun FirDeclarationDesignationWithFile.ensurePhase(firResolvePhase: FirResolvePhase, exceptTarget: Boolean) {
+            val designationIterator = toSequence(includeTarget = !exceptTarget)
+            designationIterator.forEach { firDeclaration ->
+                check(firDeclaration.resolvePhase >= firResolvePhase) {
                     "Designation element phase required to be $firResolvePhase but element resolved to ${firDeclaration.resolvePhase}"
                 }
             }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/transformers/IDEDeclarationTransformer.kt

@@ -6,49 +6,41 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.transformers
 
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.transformers.FirAbstractPhaseTransformer
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 
-internal class IDEDeclarationTransformer(private val designationIterator: FirDesignationIterator) {
+internal class IDEDeclarationTransformer(private val designation: FirDeclarationDesignation) {
+    private val designationWithoutTargetIterator = designation.toSequence(includeTarget = false).iterator()
     private var isInsideTargetDeclaration: Boolean = false
 
-    private inline fun <R> insideTargetDeclaration(insideCurrent: Boolean, action: () -> R): R {
-        val oldValue = isInsideTargetDeclaration
-        isInsideTargetDeclaration = insideCurrent
-        try {
-            return action()
-        } finally {
-            isInsideTargetDeclaration = oldValue
-        }
-    }
-
     inline fun <K, D> transformDeclarationContent(
         transformer: FirAbstractPhaseTransformer<D>,
         declaration: K,
         data: D,
-        transformDeclaration: (K, D) -> K
+        defaultCallTransform: () -> K
     ): K {
-        return if (designationIterator.canGoNext()) {
-            val declarationToTransform = designationIterator.currentDeclaration
-            val isTargetDeclaration = designationIterator.isTargetDeclaration()
-            designationIterator.goNext()
-            insideTargetDeclaration(isTargetDeclaration) {
-                declarationToTransform.visitNoTransform(transformer, data)
-            }
-            declaration
+
+        if (isInsideTargetDeclaration) {
+            return defaultCallTransform()
+        }
+
+        if (designationWithoutTargetIterator.hasNext()) {
+            designationWithoutTargetIterator.next().visitNoTransform(transformer, data)
         } else {
-            if (isInsideTargetDeclaration) {
-                transformDeclaration(declaration, data)
-            } else {
-                declaration
+            try {
+                isInsideTargetDeclaration = true
+                designation.declaration.visitNoTransform(transformer, data)
+            } finally {
+                isInsideTargetDeclaration = false
             }
         }
+
+        return declaration
     }
 
-    fun needReplacePhase(firDeclaration: FirDeclaration): Boolean =
-        isInsideTargetDeclaration
+    val needReplacePhase: Boolean get() = isInsideTargetDeclaration
 }
 
 private fun <D> FirElement.visitNoTransform(transformer: FirTransformer<D>, data: D) {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/renderer/ConeTypeIdeRenderer.kt

@@ -196,7 +196,9 @@ internal class ConeTypeIdeRenderer(
         }
 
         val classToRender = classSymbolToRender.fir
-        val designation = classToRender.collectDesignation().fullDesignation
+        val designation = classToRender.collectDesignation()
+            .toSequence(includeTarget = true)
+            .toList()
 
         var typeParametersLeft = type.typeArguments.count()
         fun needToRenderTypeParameters(index: Int): Boolean {