FIR IDE: do not fail on invalid code

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

@@ -73,6 +73,7 @@ fun KtElement.getNonLocalContainingOrThisDeclaration(): KtNamedDeclaration? {
     while (container != null && container !is KtFile) {
         if (container is KtNamedDeclaration
             && (container is KtClassOrObject || container is KtDeclarationWithBody || container is KtProperty || container is KtTypeAlias)
+            && container !is KtPrimaryConstructor
             && !KtPsiUtil.isLocal(container)
             && container !is KtEnumEntry
             && container.containingClassOrObject !is KtEnumEntry

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

@@ -5,26 +5,25 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 
-import org.jetbrains.kotlin.fir.containingClass
+import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedBodyResolveTransformerForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedBodyResolveTransformerForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedContractsResolveTransformerForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirDesignatedImplicitTypesTransformerForIDE
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtPsiUtil
 
 internal class FirLazyDeclarationResolver(
     private val firFileBuilder: FirFileBuilder
@@ -131,20 +130,14 @@ internal class FirLazyDeclarationResolver(
         val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
 
         val designation = mutableListOf<FirDeclaration>(containerFirFile)
+
         if (nonLocalDeclarationToResolve !is FirFile) {
-            val id = when (nonLocalDeclarationToResolve) {
+            val ktDeclaration = firDeclarationToResolve.ktDeclaration
-                is FirCallableDeclaration<*> -> {
+            designation += ktDeclaration.parentsOfType<KtClassOrObject>()
-                    nonLocalDeclarationToResolve.containingClass()?.classId
+                .filter { it !is KtEnumEntry }
-                }
+                .map { it.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache) }
-                is FirClassLikeDeclaration<*> -> {
+                .toList()
-                    nonLocalDeclarationToResolve.symbol.classId
+                .asReversed()
-                }
-                else -> error("Unsupported: ${nonLocalDeclarationToResolve.render()}")
-            }
-            val outerClasses = generateSequence(id) { classId ->
-                classId.outerClassId
-            }.mapTo(mutableListOf()) { provider.getFirClassifierByFqName(it)!! }
-            designation += outerClasses.asReversed()
             if (nonLocalDeclarationToResolve is FirCallableDeclaration<*>) {
                 designation += nonLocalDeclarationToResolve
             }

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

@@ -6,20 +6,31 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.util
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.psi.KtElement
 
 object FirElementFinder {
-    inline fun <reified E : FirElement> findElementIn(container: FirElement, crossinline predicate: (E) -> Boolean): E? {
+    inline fun <reified E : FirElement> findElementIn(
+        container: FirElement,
+        crossinline goInside: (E) -> Boolean = { true },
+        crossinline predicate: (E) -> Boolean,
+    ): E? {
         var result: E? = null
         container.accept(object : FirVisitorVoid() {
             override fun visitElement(element: FirElement) {
                 if (result != null) return
-                if (element is E && predicate(element)) {
+                when {
-                    result = element
+                    element !is E || element is FirFile -> {
-                } else {
+                        element.acceptChildren(this)
-                    element.acceptChildren(this)
+                    }
+                    predicate(element) -> {
+                        result = element
+                    }
+                    goInside(element) -> {
+                        element.acceptChildren(this)
+                    }
                 }
             }
         })

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

@@ -9,10 +9,12 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 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.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
@@ -22,30 +24,53 @@ internal fun KtDeclaration.findSourceNonLocalFirDeclaration(
     firSymbolProvider: FirSymbolProvider,
     moduleFileCache: ModuleFileCache
 ): FirDeclaration {
-    require(!KtPsiUtil.isLocal(this))
+    //TODO test what way faster
-    return when {
+    findSourceNonLocalFirDeclarationByProvider(firFileBuilder, firSymbolProvider, moduleFileCache)?.let { return it }
+    findSourceOfNonLocalFirDeclarationByTraversingWholeTree(firFileBuilder, moduleFileCache)?.let { return it }
+    error("No fir element was found for\n${getElementTextInContext()}")
+}
+
+private fun KtDeclaration.findSourceOfNonLocalFirDeclarationByTraversingWholeTree(
+    firFileBuilder: FirFileBuilder,
+    moduleFileCache: ModuleFileCache,
+): FirDeclaration? {
+    val firFile = firFileBuilder.buildRawFirFileWithCaching(containingKtFile, moduleFileCache, lazyBodiesMode = true)
+    val originalDeclaration = originalDeclaration
+    return FirElementFinder.findElementIn(firFile, goInside = { it is FirRegularClass }) { firDeclaration ->
+        firDeclaration.psi == this || firDeclaration.psi == originalDeclaration
+    }
+}
+
+private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider(
+    firFileBuilder: FirFileBuilder,
+    firSymbolProvider: FirSymbolProvider,
+    moduleFileCache: ModuleFileCache
+): FirDeclaration? {
+    val candidate = when {
         this is KtClassOrObject -> findFir(firSymbolProvider)
         this is KtNamedDeclaration && (this is KtProperty || this is KtNamedFunction) -> {
             val containerClass = containingClassOrObject
             val declarations = if (containerClass != null) {
                 val containerClassFir = containerClass.findFir(firSymbolProvider)
-                containerClassFir.declarations
+                containerClassFir?.declarations
             } else {
                 val ktFile = containingKtFile
                 val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache, lazyBodiesMode = true)
                 firFile.declarations
             }
             val original = originalDeclaration
-            declarations.first { it.psi === this || it.psi == original }
+            declarations?.first { it.psi === this || it.psi == original }
         }
         this is KtConstructor<*> -> {
-            val containerClassFir = containingClassOrObject?.findFir(firSymbolProvider)
+            val containingClass = containingClassOrObject
                 ?: error("Container class should be not null for KtConstructor")
+            val containerClassFir = containingClass.findFir(firSymbolProvider) ?: return null
             containerClassFir.declarations.first { it.psi === this }
         }
         this is KtTypeAlias -> findFir(firSymbolProvider)
         else -> error("Invalid container $this::class")
     }
+    return candidate?.takeIf { it.realPsi == this }
 }
 
 val ORIGINAL_DECLARATION_KEY = com.intellij.openapi.util.Key<KtDeclaration>("ORIGINAL_DECLARATION_KEY")
@@ -53,9 +78,8 @@ val ORIGINAL_DECLARATION_KEY = com.intellij.openapi.util.Key<KtDeclaration>("ORI
 var KtDeclaration.originalDeclaration by UserDataProperty(ORIGINAL_DECLARATION_KEY)
 
 
-private fun KtClassOrObject.findFir(firSymbolProvider: FirSymbolProvider): FirRegularClass {
+private fun KtClassOrObject.findFir(firSymbolProvider: FirSymbolProvider): FirRegularClass? {
-    val classId = classIdIfNonLocal()
+    val classId = classIdIfNonLocal() ?: return null
-        ?: error("Container classId should not be null for non-local declaration")
     return executeWithoutPCE {
         firSymbolProvider.getClassLikeSymbolByFqName(classId)?.fir as? FirRegularClass
             ?: error("Could not find class $classId")

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

@@ -8,8 +8,12 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.util
 import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.openapi.progress.ProgressManager
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import org.jetbrains.kotlin.psi.KtDeclaration
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.locks.ReentrantLock
 
@@ -52,6 +56,13 @@ internal inline fun checkCanceled() {
 internal val FirElement.isErrorElement
     get() = this is FirDiagnosticHolder
 
+internal val FirDeclaration.ktDeclaration: KtDeclaration
+    get() {
+        val psi = psi
+            ?: error("PSI element was not found for${render()}")
+        return psi as KtDeclaration
+    }
+
 
 internal fun IdeaModuleInfo.collectTransitiveDependenciesWithSelf(): List<IdeaModuleInfo> {
     val result = mutableSetOf<IdeaModuleInfo>()