[LL FIR] make FirDesignation.path type more precise. I's always a FirRegularClass

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

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLookupTagWithFixedSymbol
 
 class FirDesignationWithFile(
-    path: List<FirDeclaration>,
+    path: List<FirRegularClass>,
     target: FirElementWithResolvePhase,
     val firFile: FirFile
 ) : FirDesignation(
@@ -40,7 +40,7 @@ class FirDesignationWithFile(
 }
 
 open class FirDesignation(
-    val path: List<FirDeclaration>,
+    val path: List<FirRegularClass>,
     val target: FirElementWithResolvePhase,
 ) {
     val firstNonFileDeclaration: FirElementWithResolvePhase
@@ -52,20 +52,20 @@ open class FirDesignation(
     }
 }
 
-private fun FirRegularClass.collectForNonLocal(): List<FirDeclaration> {
+private fun FirRegularClass.collectForNonLocal(): List<FirRegularClass> {
     require(!isLocal)
     val firProvider = moduleData.session.firProvider
     var containingClassId = classId.outerClassId
-    val designation = mutableListOf<FirDeclaration>(this)
+    val designation = mutableListOf<FirRegularClass>(this)
     while (containingClassId != null) {
-        val currentClass = firProvider.getFirClassifierByFqName(containingClassId) ?: break
+        val currentClass = firProvider.getFirClassifierByFqName(containingClassId) as? FirRegularClass ?: break
         designation.add(currentClass)
         containingClassId = containingClassId.outerClassId
     }
     return designation
 }
 
-private fun collectDesignationPath(target: FirElementWithResolvePhase): List<FirDeclaration>? {
+private fun collectDesignationPath(target: FirElementWithResolvePhase): List<FirRegularClass>? {
     val containingClass = when (target) {
         is FirCallableDeclaration -> {
             if (target !is FirConstructor && target.symbol.callableId.isLocal) return null

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/LLFirDesignatedBodyResolveTransformerForReturnTypeCalculator.kt

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
+import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
@@ -17,8 +19,6 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyRe
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator
 
 fun LLFirDesignatedImpliciteTypesBodyResolveTransformerForReturnTypeCalculator(
     designation: Iterator<FirElement>,
@@ -67,7 +67,10 @@ private class LLFirDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl(
         if (returnTypeRef is FirImplicitTypeRef) {
             val declarationList = designation.filterIsInstance<FirDeclaration>()
             check(declarationList.isNotEmpty()) { "Invalid empty declaration designation" }
-            body(FirDesignation(declarationList.dropLast(1), this))
+            val designationPath = declarationList.dropLast(1)
+            check(designationPath.all { it is FirRegularClass })
+            @Suppress("UNCHECKED_CAST")
+            body(FirDesignation(designationPath as List<FirRegularClass>, this))
         }
     }
 

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

@@ -26,14 +26,14 @@ import org.jetbrains.kotlin.psi.*
 
 internal object FirLazyBodiesCalculator {
     fun calculateLazyBodiesInside(designation: FirDesignation) {
-        designation.target.transform<FirElement, PersistentList<FirDeclaration>>(
+        designation.target.transform<FirElement, PersistentList<FirRegularClass>>(
             FirLazyBodiesCalculatorTransformer,
             designation.path.toPersistentList(),
         )
     }
 
     fun calculateLazyBodies(firFile: FirFile) {
-        firFile.transform<FirElement, PersistentList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, persistentListOf())
+        firFile.transform<FirElement, PersistentList<FirRegularClass>>(FirLazyBodiesCalculatorTransformer, persistentListOf())
     }
 
     private fun replaceValueParameterDefaultValues(valueParameters: List<FirValueParameter>, newValueParameters: List<FirValueParameter>) {
@@ -171,16 +171,16 @@ internal object FirLazyBodiesCalculator {
                 || firProperty.getExplicitBackingField()?.initializer is FirLazyExpression
 }
 
-private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentList<FirDeclaration>>() {
+private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentList<FirRegularClass>>() {
 
-    override fun transformFile(file: FirFile, data: PersistentList<FirDeclaration>): FirFile {
+    override fun transformFile(file: FirFile, data: PersistentList<FirRegularClass>): FirFile {
         file.declarations.forEach {
             it.transformSingle(this, data)
         }
         return file
     }
 
-    override fun <E : FirElement> transformElement(element: E, data: PersistentList<FirDeclaration>): E {
+    override fun <E : FirElement> transformElement(element: E, data: PersistentList<FirRegularClass>): E {
         if (element is FirRegularClass) {
             val newList = data.add(element)
             element.declarations.forEach {
@@ -193,7 +193,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
 
     override fun transformSimpleFunction(
         simpleFunction: FirSimpleFunction,
-        data: PersistentList<FirDeclaration>
+        data: PersistentList<FirRegularClass>
     ): FirSimpleFunction {
         if (FirLazyBodiesCalculator.needCalculatingLazyBodyForFunction(simpleFunction)) {
             val designation = FirDesignation(data, simpleFunction)
@@ -204,7 +204,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
 
     override fun transformConstructor(
         constructor: FirConstructor,
-        data: PersistentList<FirDeclaration>
+        data: PersistentList<FirRegularClass>
     ): FirConstructor {
         if (FirLazyBodiesCalculator.needCalculatingLazyBodyForConstructor(constructor)) {
             val designation = FirDesignation(data, constructor)
@@ -213,7 +213,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
         return constructor
     }
 
-    override fun transformProperty(property: FirProperty, data: PersistentList<FirDeclaration>): FirProperty {
+    override fun transformProperty(property: FirProperty, data: PersistentList<FirRegularClass>): FirProperty {
         if (FirLazyBodiesCalculator.needCalculatingLazyBodyForProperty(property)) {
             val designation = FirDesignation(data, property)
             FirLazyBodiesCalculator.calculateLazyBodyForProperty(designation)
@@ -221,11 +221,11 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
         return property
     }
 
-    override fun transformPropertyAccessor(propertyAccessor: FirPropertyAccessor, data: PersistentList<FirDeclaration>): FirStatement {
+    override fun transformPropertyAccessor(propertyAccessor: FirPropertyAccessor, data: PersistentList<FirRegularClass>): FirStatement {
         return propertyAccessor.also { transformProperty(it.propertySymbol.fir, data) }
     }
 
-    override fun transformEnumEntry(enumEntry: FirEnumEntry, data: PersistentList<FirDeclaration>): FirStatement {
+    override fun transformEnumEntry(enumEntry: FirEnumEntry, data: PersistentList<FirRegularClass>): FirStatement {
         if (enumEntry.initializer is FirLazyExpression) {
             val designation = FirDesignation(data, enumEntry)
             FirLazyBodiesCalculator.calculateLazyInitializerForEnumEntry(designation)
@@ -234,7 +234,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
     }
 
     override fun transformAnonymousInitializer(
-        anonymousInitializer: FirAnonymousInitializer, data: PersistentList<FirDeclaration>
+        anonymousInitializer: FirAnonymousInitializer, data: PersistentList<FirRegularClass>
     ): FirAnonymousInitializer {
         if (anonymousInitializer.body is FirLazyBlock) {
             val designation = FirDesignation(data, anonymousInitializer)

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

@@ -21,7 +21,7 @@ import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 
 /**
- * Transform designation into STATUS phase. Affects only for designation, target declaration, it's children and dependents
+ * Transform designation into STATUS phase. Affects only for designation, members of designation classes, target declaration, it's children and dependents
  */
 internal class LLFirDesignatedStatusResolveTransformer(
     private val designation: FirDesignationWithFile,
@@ -29,12 +29,11 @@ internal class LLFirDesignatedStatusResolveTransformer(
     private val scopeSession: ScopeSession,
 ) : LLFirLazyTransformer {
     private inner class FirDesignatedStatusResolveTransformerForIDE(
-        private val designationIterator: Iterator<FirElementWithResolvePhase>
+        private val designationIterator: Iterator<FirRegularClass>,
     ) : FirStatusResolveTransformer(session, scopeSession, StatusComputationSession()) {
 
         private var isInsideTargetDeclaration = false
 
-
         override fun transformClass(klass: FirClass, data: FirResolvedDeclarationStatus?): FirStatement {
             return storeClass(klass) {
                 resolveClassTypeParameters(klass)
@@ -71,14 +70,8 @@ internal class LLFirDesignatedStatusResolveTransformer(
                 designation.target.transform<FirDeclaration, _>(this, data = null)
                 return
             }
-            when (val nextElement = designationIterator.next()) {
-                is FirClass -> {
-                    transformClassContent(nextElement, data = null)
-                }
-                else -> {
-                    error("Unexpected declaration in designation: ${nextElement::class.qualifiedName}")
-                }
-            }
+            val nextClass = designationIterator.next()
+            transformClassContent(nextClass, data = null)
         }
     }
 
@@ -87,7 +80,7 @@ internal class LLFirDesignatedStatusResolveTransformer(
         if (designation.target.resolvePhase >= FirResolvePhase.STATUS) return
         designation.target.checkPhase(FirResolvePhase.TYPES)
 
-        val designationIterator = designation.toSequence(includeTarget = false).iterator()
+        val designationIterator = designation.path.iterator()
         val transformer = FirDesignatedStatusResolveTransformerForIDE(designationIterator)
         phaseRunner.runPhaseWithCustomResolve(FirResolvePhase.STATUS) {
             transformer.moveNextDeclaration()

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

@@ -137,7 +137,6 @@ internal class LLFirDesignatedSupertypeResolverTransformer(
                 updatePhaseDeep(designation.target, FirResolvePhase.SUPER_TYPES)
                 return
             }
-            check(resolvableTarget is FirClassLikeDeclaration)
             val targetPath = designation.path.dropLast(1)
             FirDesignationWithFile(targetPath, resolvableTarget, designation.firFile)
         } else designation

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractPartialRawFirBuilderTestCase.kt

@@ -58,7 +58,7 @@ abstract class AbstractPartialRawFirBuilderTestCase : AbstractLowLevelApiSingleF
     }
 
     private class DesignationBuilder(private val elementToBuild: KtDeclaration) : FirVisitorVoid() {
-        private val path = mutableListOf<FirDeclaration>()
+        private val path = mutableListOf<FirRegularClass>()
         var resultDesignation: FirDesignation? = null
             private set