FIR IDE: do not fail on invalid code

2020-10-20 17:04:04 +03:00
parent c646cb3d7e
commit 7dcda00e1d
5 changed files with 72 additions and 32 deletions
@@ -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
@@ -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.KtDeclaration
-import org.jetbrains.kotlin.psi.KtPsiUtil
+import org.jetbrains.kotlin.psi.*

 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) {
-                is FirCallableDeclaration<*> -> {
-                    nonLocalDeclarationToResolve.containingClass()?.classId
-                }
-                is FirClassLikeDeclaration<*> -> {
-                    nonLocalDeclarationToResolve.symbol.classId
-                }
-                else -> error("Unsupported: ${nonLocalDeclarationToResolve.render()}")
-            }
-            val outerClasses = generateSequence(id) { classId ->
-                classId.outerClassId
-            }.mapTo(mutableListOf()) { provider.getFirClassifierByFqName(it)!! }
-            designation += outerClasses.asReversed()
+            val ktDeclaration = firDeclarationToResolve.ktDeclaration
+            designation += ktDeclaration.parentsOfType<KtClassOrObject>()
+                .filter { it !is KtEnumEntry }
+                .map { it.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache) }
+                .toList()
+                .asReversed()
            if (nonLocalDeclarationToResolve is FirCallableDeclaration<*>) {
                designation += nonLocalDeclarationToResolve
            }
@@ -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)) {
-                    result = element
-                } else {
-                    element.acceptChildren(this)
+                when {
+                    element !is E || element is FirFile -> {
+                        element.acceptChildren(this)
+                    }
+                    predicate(element) -> {
+                        result = element
+                    }
+                    goInside(element) -> {
+                        element.acceptChildren(this)
+                    }
                }
            }
        })
@@ -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))
-    return when {
+    //TODO test what way faster
+    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 {
-    val classId = classIdIfNonLocal()
-        ?: error("Container classId should not be null for non-local declaration")
+private fun KtClassOrObject.findFir(firSymbolProvider: FirSymbolProvider): FirRegularClass? {
+    val classId = classIdIfNonLocal() ?: return null
    return executeWithoutPCE {
        firSymbolProvider.getClassLikeSymbolByFqName(classId)?.fir as? FirRegularClass
            ?: error("Could not find class $classId")
@@ -68,4 +92,4 @@ private fun KtTypeAlias.findFir(firSymbolProvider: FirSymbolProvider): FirTypeAl
        firSymbolProvider.getClassLikeSymbolByFqName(typeAlias)?.fir as? FirTypeAlias
            ?: error("Could not find type alias $typeAlias")
    }
-}
+}
@@ -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>()