[FIR IDE] Remove designation collecting for local declarations

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

@@ -8,19 +8,21 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.api
 import org.jetbrains.kotlin.fir.containingClass
 import org.jetbrains.kotlin.fir.containingClassForLocal
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
+import org.jetbrains.kotlin.idea.util.ifFalse
+import org.jetbrains.kotlin.idea.util.ifTrue
 
 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)
@@ -32,7 +34,6 @@ class FirDeclarationDesignationWithFile(
 open class FirDeclarationDesignation(
     val path: List<FirDeclaration>,
     val declaration: FirDeclaration,
-    val isLocalDesignation: Boolean
 ) {
     fun toSequence(includeTarget: Boolean): Sequence<FirDeclaration> = sequence {
         yieldAll(path)
@@ -40,29 +41,6 @@ open class FirDeclarationDesignation(
     }
 }
 
-private fun FirClassLikeDeclaration<*>.containingClass(): FirClassLikeDeclaration<*>? =
-    if (isLocal) (this as? FirRegularClass)?.containingClassForLocal()?.toFirRegularClass(moduleData.session)
-    else symbol.classId.outerClassId?.let(moduleData.session.firProvider::getFirClassifierByFqName)
-
-private fun collectDesignationAndIsLocal(declaration: FirDeclaration): Pair<List<FirDeclaration>, Boolean> {
-    val containingClass = when (declaration) {
-        is FirCallableDeclaration<*> -> declaration.containingClass()?.toFirRegularClass(declaration.moduleData.session)
-        is FirAnonymousObject -> return emptyList<FirDeclaration>() to true
-        is FirClassLikeDeclaration<*> -> declaration.containingClass()
-        else -> error("Invalid declaration ${declaration.renderWithType()}")
-    } ?: return emptyList<FirDeclaration>() to 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 path.reversed() to containingClass.isLocal
-}
-
 private fun FirRegularClass.collectForNonLocal(): List<FirDeclaration> {
     require(!isLocal)
     val firProvider = moduleData.session.firProvider
@@ -76,25 +54,48 @@ private fun FirRegularClass.collectForNonLocal(): List<FirDeclaration> {
     return designation
 }
 
-private fun FirRegularClass.collectForLocal(): List<FirClassLikeDeclaration<*>> {
+private fun collectDesignationPath(declaration: FirDeclaration): List<FirDeclaration>? {
-    require(isLocal)
+    val containingClass = when (declaration) {
-    var containingClassLookUp = containingClassForLocal()
+        is FirCallableDeclaration<*> ->
-    val designation = mutableListOf<FirClassLikeDeclaration<*>>(this)
+            when (declaration) {
-    while (containingClassLookUp != null && containingClassLookUp.classId.isLocal) {
+                is FirSimpleFunction, is FirProperty, is FirField, is FirConstructor ->
-        val currentClass = containingClassLookUp.toFirRegularClass(moduleData.session) ?: break
+                    declaration.containingClass()?.toFirRegularClass(declaration.moduleData.session)
-        designation.add(currentClass)
+                else -> return null
-        containingClassLookUp = currentClass.containingClassForLocal()
+            }
+        is FirClassLikeDeclaration<*> -> {
+            if (declaration.isLocal) return null
+            declaration.symbol.classId.outerClassId?.let(declaration.moduleData.session.firProvider::getFirClassifierByFqName)
+        }
+        else -> return null
+    } ?: return emptyList()
+
+    require(containingClass is FirRegularClass) {
+        "FirRegularClass as containing declaration expected but found ${containingClass.renderWithType()}"
     }
-    return designation
+    return containingClass.isLocal.ifFalse { containingClass.collectForNonLocal().asReversed() }
 }
 
-
 fun FirDeclaration.collectDesignation(firFile: FirFile): FirDeclarationDesignationWithFile =
-    collectDesignationAndIsLocal(this).let {
+    tryCollectDesignation(firFile) ?: error("No designation of local declaration ${this.render()}")
-        FirDeclarationDesignationWithFile(it.first, this, it.second, firFile)
-    }
 
 fun FirDeclaration.collectDesignation(): FirDeclarationDesignation =
-    collectDesignationAndIsLocal(this).let {
+    tryCollectDesignation() ?: error("No designation of local declaration ${this.render()}")
-        FirDeclarationDesignation(it.first, this, it.second)
+
+fun FirDeclaration.collectDesignationWithFile(): FirDeclarationDesignationWithFile =
+    tryCollectDesignationWithFile() ?: error("No designation of local declaration ${this.render()}")
+
+fun FirDeclaration.tryCollectDesignation(firFile: FirFile): FirDeclarationDesignationWithFile? =
+    collectDesignationPath(this)?.let {
+        FirDeclarationDesignationWithFile(it, this, firFile)
     }
+
+fun FirDeclaration.tryCollectDesignation(): FirDeclarationDesignation? =
+    collectDesignationPath(this)?.let {
+        FirDeclarationDesignation(it, this)
+    }
+
+fun FirDeclaration.tryCollectDesignationWithFile(): FirDeclarationDesignationWithFile? {
+    val path = collectDesignationPath(this) ?: return null
+    val firFile = getContainingFile() ?: return null
+    return FirDeclarationDesignationWithFile(path, this, firFile)
+}

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

@@ -268,7 +268,6 @@ object LowLevelFirApiFacadeForResolveOnAir {
             val onAirDesignation = FirDeclarationDesignationWithFile(
                 path = originalDesignation.path,
                 declaration = copiedFirDeclaration,
-                isLocalDesignation = false,
                 firFile = originalFirFile
             )
             state.firLazyDeclarationResolver.runLazyDesignatedOnAirResolveToBodyWithoutLock(

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

@@ -63,9 +63,6 @@ internal object PersistenceContextCollector {
             "Cannot collect context for local declaration ${declaration.renderWithType()}"
         }
         val designation = declaration.collectDesignation(firFile)
-        check(!designation.isLocalDesignation) {
-            "Designation should not local for ${declaration.renderWithType()}"
-        }
         return ContextCollectingDiagnosticCollectorVisitor.collect(sessionHolder, designation)
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator.kt

@@ -62,19 +62,19 @@ class FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl(
         simpleFunction: FirSimpleFunction,
         data: ResolutionMode
     ): FirSimpleFunction {
-        val firDesignation = FirDeclarationDesignation(declarationDesignation, simpleFunction, false)
+        val firDesignation = FirDeclarationDesignation(declarationDesignation, simpleFunction)
         FirLazyBodiesCalculator.calculateLazyBodiesForFunction(firDesignation)
         return super.transformSimpleFunction(simpleFunction, data)
     }
 
     override fun transformConstructor(constructor: FirConstructor, data: ResolutionMode): FirDeclaration {
-        val firDesignation = FirDeclarationDesignation(declarationDesignation, constructor, false)
+        val firDesignation = FirDeclarationDesignation(declarationDesignation, constructor)
         FirLazyBodiesCalculator.calculateLazyBodyForSecondaryConstructor(firDesignation)
         return super.transformConstructor(constructor, data)
     }
 
     override fun transformProperty(property: FirProperty, data: ResolutionMode): FirProperty {
-        val firDesignation = FirDeclarationDesignation(declarationDesignation, property, false)
+        val firDesignation = FirDeclarationDesignation(declarationDesignation, property)
         FirLazyBodiesCalculator.calculateLazyBodyForProperty(firDesignation)
         return super.transformProperty(property, data)
     }

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

@@ -47,7 +47,6 @@ internal object FirLazyBodiesCalculator {
     }
 
     fun calculateLazyBodyForSecondaryConstructor(designation: FirDeclarationDesignation) {
-        require(!designation.isLocalDesignation) { "Not supported for local designations" }
         val secondaryConstructor = designation.declaration as FirConstructor
         require(!secondaryConstructor.isPrimary)
         if (secondaryConstructor.body !is FirLazyBlock) return
@@ -65,7 +64,6 @@ internal object FirLazyBodiesCalculator {
     }
 
     fun calculateLazyBodyForProperty(designation: FirDeclarationDesignation) {
-        require(!designation.isLocalDesignation) { "Not supported for local designations" }
         val firProperty = designation.declaration as FirProperty
         if (!needCalculatingLazyBodyForProperty(firProperty)) return
 
@@ -132,7 +130,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<F
         data: MutableList<FirDeclaration>
     ): FirDeclaration {
         if (simpleFunction.body is FirLazyBlock) {
-            val designation = FirDeclarationDesignation(data, simpleFunction, false)
+            val designation = FirDeclarationDesignation(data, simpleFunction)
             FirLazyBodiesCalculator.calculateLazyBodiesForFunction(designation)
         }
         return simpleFunction
@@ -143,7 +141,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<F
         data: MutableList<FirDeclaration>
     ): FirDeclaration {
         if (constructor.body is FirLazyBlock) {
-            val designation = FirDeclarationDesignation(data, constructor, false)
+            val designation = FirDeclarationDesignation(data, constructor)
             FirLazyBodiesCalculator.calculateLazyBodyForSecondaryConstructor(designation)
         }
         return constructor
@@ -151,7 +149,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<F
 
     override fun transformProperty(property: FirProperty, data: MutableList<FirDeclaration>): FirDeclaration {
         if (FirLazyBodiesCalculator.needCalculatingLazyBodyForProperty(property)) {
-            val designation = FirDeclarationDesignation(data, property, false)
+            val designation = FirDeclarationDesignation(data, property)
             FirLazyBodiesCalculator.calculateLazyBodyForProperty(designation)
         }
         return property

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

@@ -251,8 +251,6 @@ internal class FirLazyDeclarationResolver(
         checkPCE: Boolean,
         declarationPhaseDowngraded: Boolean,
     ) {
-        check(!designation.isLocalDesignation) { "Could not resolve local designation" }
-
         val filePhase = designation.firFile.resolvePhase
         if (filePhase == FirResolvePhase.RAW_FIR) {
             lazyResolveFileDeclarationWithoutLock(

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

@@ -36,8 +36,6 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
             rootNonLocalDeclaration: KtDeclaration,
             replacement: RawFirReplacement?
         ): FirDeclaration {
-            require(!designation.isLocalDesignation) { "Local declaration is not allowed here" }
-
             val replacementApplier = replacement?.Applier()
             val builder = RawFirNonLocalDeclarationBuilder(
                 session = session,
@@ -57,8 +55,6 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
             designation: FirDeclarationDesignation,
             rootNonLocalDeclaration: KtDeclaration,
         ): FirDeclaration {
-            require(!designation.isLocalDesignation) { "Local declaration is not allowed here" }
-
             val functionsToRebind = when (val originalDeclaration = designation.declaration) {
                 is FirSimpleFunction -> setOf(originalDeclaration)
                 is FirProperty -> setOfNotNull(originalDeclaration.getter, originalDeclaration.setter)

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

@@ -159,7 +159,7 @@ internal class FirDesignatedSupertypeResolverTransformerForIDE(
             val resolvableTarget = designation.path.lastOrNull() ?: return
             check(resolvableTarget is FirClassLikeDeclaration<*>)
             val targetPath = designation.path.dropLast(1)
-            FirDeclarationDesignationWithFile(targetPath, resolvableTarget, false, designation.firFile)
+            FirDeclarationDesignationWithFile(targetPath, resolvableTarget, designation.firFile)
         } else designation
 
         if (targetDesignation.isResolvedForAllDeclarations(FirResolvePhase.SUPER_TYPES, declarationPhaseDowngraded)) return

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

@@ -5,8 +5,12 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.util
 
-import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.InvalidFirElementTypeException
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration

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

@@ -6,20 +6,22 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.renderer
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.containingClass
-import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
+import org.jetbrains.kotlin.fir.containingClassForLocal
-import org.jetbrains.kotlin.fir.declarations.isInner
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.toAnnotationClassId
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedError
+import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.inference.*
+import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.asJava.applyIf
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.tryCollectDesignation
 import org.jetbrains.kotlin.idea.frontend.api.components.KtTypeRendererOptions
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.idea.util.ifTrue
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.renderer.render
 
@@ -176,6 +178,32 @@ internal class ConeTypeIdeRenderer(
         return null
     }
 
+    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
+    }
+
+    private fun collectDesignationPathForLocal(declaration: FirDeclaration): List<FirDeclaration>? {
+        val containingClass = when (declaration) {
+            is FirCallableDeclaration<*> -> declaration.containingClass()?.toFirRegularClass(declaration.moduleData.session)
+            is FirAnonymousObject -> return listOf(declaration)
+            is FirClassLikeDeclaration<*> -> declaration.let {
+                if (!declaration.isLocal) return null
+                (it as? FirRegularClass)?.containingClassForLocal()?.toFirRegularClass(declaration.moduleData.session)
+            }
+            else -> error("Invalid declaration ${declaration.renderWithType()}")
+        } ?: return listOf(declaration)
+
+        return containingClass.isLocal.ifTrue { containingClass.collectForLocal().reversed() }
+    }
+
     private fun StringBuilder.renderTypeConstructorAndArguments(type: ConeClassLikeType) {
         fun renderTypeArguments(typeArguments: Array<out ConeTypeProjection>, range: IntRange) {
             if (range.any()) {
@@ -208,10 +236,12 @@ internal class ConeTypeIdeRenderer(
             return
         }
 
-        val classToRender = classSymbolToRender.fir
+        val designation = classSymbolToRender.fir.let {
-        val designation = classToRender.collectDesignation()
+            val nonLocalDesignation = it.tryCollectDesignation()
-            .toSequence(includeTarget = true)
+            nonLocalDesignation?.toSequence(includeTarget = true)?.toList()
-            .toList()
+                ?: collectDesignationPathForLocal(it)
+                ?: emptyList()
+        }
 
         var typeParametersLeft = type.typeArguments.count()
         fun needToRenderTypeParameters(index: Int): Boolean {