LL API: do not lazy resolve local declarations

2022-01-20 19:02:11 +01:00
parent c7729689c4
commit 6dd01a2ff5
7 changed files with 68 additions and 115 deletions
@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.KtToFirMapping
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.ResolveType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.containingKtFileIfAny
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalKtFile
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
@@ -62,11 +61,9 @@ internal class FirModuleResolveStateDepended(
    override fun getOrBuildFirFile(ktFile: KtFile): FirFile =
        originalState.getOrBuildFirFile(ktFile)

-    override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D =
+    override fun resolveFirToPhase(declaration: FirDeclaration, toPhase: FirResolvePhase) {
        originalState.resolveFirToPhase(declaration, toPhase)
-
-    override fun <D : FirDeclaration> resolveFirToResolveType(declaration: D, type: ResolveType): D =
-        originalState.resolveFirToResolveType(declaration, type)
+    }

    override fun tryGetCachedFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? {
        val ktFile = declaration.containingKtFileIfAny ?: return null
@@ -16,8 +16,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.FirFileBuild
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FileStructureCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyDeclarationResolver
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.ResolveType
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.lazyResolveDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.firIdeProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSourcesSession
@@ -28,9 +26,9 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
@@ -73,7 +71,7 @@ internal class FirModuleResolveStateImpl(
        )

    override fun getOrBuildFirFile(ktFile: KtFile): FirFile =
-        firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache, preferLazyBodies = false)
+        firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache)

    override fun tryGetCachedFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? =
        cache.getContainerFirFile(declaration)
@@ -139,13 +137,13 @@ internal class FirModuleResolveStateImpl(
        }
    }

-    override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
-        if (toPhase == FirResolvePhase.RAW_FIR) return declaration
+    override fun resolveFirToPhase(declaration: FirDeclaration, toPhase: FirResolvePhase) {
+        if (toPhase == FirResolvePhase.RAW_FIR) return
        val fileCache = when (val session = declaration.moduleData.session) {
            is FirIdeSourcesSession -> session.cache
-            else -> return declaration
+            else -> return
        }
-        return firLazyDeclarationResolver.lazyResolveDeclaration(
+        firLazyDeclarationResolver. lazyResolveDeclaration(
            firDeclarationToResolve = declaration,
            moduleFileCache = fileCache,
            scopeSession = ScopeSession(),
@@ -153,19 +151,4 @@ internal class FirModuleResolveStateImpl(
            checkPCE = true,
        )
    }
-
-    override fun <D : FirDeclaration> resolveFirToResolveType(declaration: D, type: ResolveType): D {
-        if (type == ResolveType.NoResolve) return declaration
-        val fileCache = when (val session = declaration.moduleData.session) {
-            is FirIdeSourcesSession -> session.cache
-            else -> return declaration
-        }
-        return firLazyDeclarationResolver.lazyResolveDeclaration(
-            firDeclaration = declaration,
-            moduleFileCache = fileCache,
-            scopeSession = ScopeSession(),
-            toResolveType = type,
-            checkPCE = true,
-        )
-    }
 }
@@ -33,7 +33,6 @@ internal class IdeFirPhaseManager(
                scopeSession = ScopeSession(),
                toPhase = requiredPhase,
                checkPCE = true,
-                skipLocalDeclaration = true,
            )
        } catch (e: Throwable) {
            sessionInvalidator.invalidate(fir.moduleData.session)
@@ -116,7 +116,7 @@ internal class FileStructure private constructor(
            moduleFileCache,
            firFile
        )
-        val resolvedDeclaration = firLazyDeclarationResolver.lazyResolveDeclaration(
+        firLazyDeclarationResolver.lazyResolveDeclaration(
            firDeclarationToResolve = firDeclaration,
            moduleFileCache = moduleFileCache,
            scopeSession = ScopeSession(),
@@ -124,7 +124,7 @@ internal class FileStructure private constructor(
            checkPCE = true,
        )
        return FileElementFactory.createFileStructureElement(
-            firDeclaration = resolvedDeclaration,
+            firDeclaration = firDeclaration,
            ktDeclaration = declaration,
            firFile = firFile,
            firFileLockProvider = moduleFileCache.firFileLockProvider
@@ -133,7 +133,7 @@ internal class FileStructure private constructor(

    private fun createStructureElement(container: KtAnnotated): FileStructureElement = when (container) {
        is KtFile -> {
-            val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache, preferLazyBodies = true)
+            val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
            firLazyDeclarationResolver.resolveFileAnnotations(
                firFile = firFile,
                annotations = firFile.annotations,
@@ -140,16 +140,13 @@ internal class ReanalyzableFunctionStructureElement(
                it.replaceResolvePhase(minOf(it.resolvePhase, upgradedPhase))
            }

-            val resolvedDeclaration = firLazyDeclarationResolver.lazyResolveDeclaration(
+            firLazyDeclarationResolver.lazyResolveDeclaration(
                firDeclarationToResolve = originalFunction,
                moduleFileCache = cache,
                scopeSession = ScopeSession(),
                toPhase = FirResolvePhase.BODY_RESOLVE,
                checkPCE = true,
            )
-            check(resolvedDeclaration === originalFunction) {
-                "Reanalysed declaration not expected to be updated"
-            }

            ReanalyzableFunctionStructureElement(
                firFile,
@@ -203,16 +200,14 @@ internal class ReanalyzablePropertyStructureElement(
                replaceBodyResolveState(FirPropertyBodyResolveState.NOTHING_RESOLVED)
            }

-            val resolvedDeclaration = firLazyDeclarationResolver.lazyResolveDeclaration(
+            firLazyDeclarationResolver.lazyResolveDeclaration(
                firDeclarationToResolve = originalProperty,
                moduleFileCache = cache,
                scopeSession = ScopeSession(),
                toPhase = FirResolvePhase.BODY_RESOLVE,
                checkPCE = true,
            )
-            check(resolvedDeclaration === originalProperty) {
-                "Reanalysed declaration not expected to be updated"
-            }
+

            ReanalyzablePropertyStructureElement(
                firFile,
@@ -5,6 +5,9 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve

+import kotlinx.collections.immutable.PersistentList
+import kotlinx.collections.immutable.persistentListOf
+import kotlinx.collections.immutable.toPersistentList
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.firIdeProvider
 import org.jetbrains.kotlin.fir.FirElement
@@ -23,14 +26,14 @@ import org.jetbrains.kotlin.psi.KtSecondaryConstructor

 internal object FirLazyBodiesCalculator {
    fun calculateLazyBodiesInside(designation: FirDeclarationDesignation) {
-        designation.declaration.transform<FirElement, MutableList<FirDeclaration>>(
+        designation.declaration.transform<FirElement, PersistentList<FirDeclaration>>(
            FirLazyBodiesCalculatorTransformer,
-            designation.toSequence(includeTarget = true).toMutableList()
+            designation.toSequence(includeTarget = false).toList().toPersistentList()
        )
    }

    fun calculateLazyBodies(firFile: FirFile) {
-        firFile.transform<FirElement, MutableList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, mutableListOf())
+        firFile.transform<FirElement, PersistentList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, persistentListOf())
    }

    fun calculateLazyBodiesForFunction(designation: FirDeclarationDesignation) {
@@ -122,30 +125,29 @@ internal object FirLazyBodiesCalculator {
                || firProperty.getExplicitBackingField()?.initializer is FirLazyExpression
 }

-private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<FirDeclaration>>() {
+private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentList<FirDeclaration>>() {

-    override fun transformFile(file: FirFile, data: MutableList<FirDeclaration>): FirFile {
+    override fun transformFile(file: FirFile, data: PersistentList<FirDeclaration>): FirFile {
        file.declarations.forEach {
            it.transformSingle(this, data)
        }
        return file
    }

-    override fun <E : FirElement> transformElement(element: E, data: MutableList<FirDeclaration>): E {
+    override fun <E : FirElement> transformElement(element: E, data: PersistentList<FirDeclaration>): E {
        if (element is FirRegularClass) {
-            data.add(element)
+            val newList = data.add(element)
            element.declarations.forEach {
-                it.transformSingle(this, data)
+                it.transformSingle(this, newList)
            }
-            element.transformChildren(this, data)
-            data.removeLast()
+            element.transformChildren(this, newList)
        }
        return element
    }

    override fun transformSimpleFunction(
        simpleFunction: FirSimpleFunction,
-        data: MutableList<FirDeclaration>
+        data: PersistentList<FirDeclaration>
    ): FirSimpleFunction {
        if (simpleFunction.body is FirLazyBlock) {
            val designation = FirDeclarationDesignation(data, simpleFunction)
@@ -156,7 +158,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<F

    override fun transformConstructor(
        constructor: FirConstructor,
-        data: MutableList<FirDeclaration>
+        data: PersistentList<FirDeclaration>
    ): FirConstructor {
        if (constructor.body is FirLazyBlock) {
            val designation = FirDeclarationDesignation(data, constructor)
@@ -165,7 +167,7 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<F
        return constructor
    }

-    override fun transformProperty(property: FirProperty, data: MutableList<FirDeclaration>): FirProperty {
+    override fun transformProperty(property: FirProperty, data: PersistentList<FirDeclaration>): FirProperty {
        if (FirLazyBodiesCalculator.needCalculatingLazyBodyForProperty(property)) {
            val designation = FirDeclarationDesignation(data, property)
            FirLazyBodiesCalculator.calculateLazyBodyForProperty(designation)
@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.firIdeProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.FirFileAnnotationsResolveTransformer
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.FirProviderInterceptorForIDE
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LazyTransformerFactory
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.fir.FirElement
@@ -24,15 +23,13 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.FirImportResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
-import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtClassBody
 import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtEnumEntry

 internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBuilder) {
    /**
@@ -227,18 +224,17 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
     * @see lazyResolveDeclaration for ordinary resolve
     * @param firDeclarationToResolve target non-local declaration
     */
-    fun <D : FirDeclaration> lazyResolveDeclaration(
-        firDeclarationToResolve: D,
+    fun lazyResolveDeclaration(
+        firDeclarationToResolve: FirDeclaration,
        moduleFileCache: ModuleFileCache,
        scopeSession: ScopeSession,
        toPhase: FirResolvePhase,
        checkPCE: Boolean,
-        skipLocalDeclaration: Boolean = false,
-    ): D {
-        if (toPhase == FirResolvePhase.RAW_FIR) return firDeclarationToResolve
+    ) {
+        if (toPhase == FirResolvePhase.RAW_FIR) return
        if (toPhase == FirResolvePhase.IMPORTS) {
            if (fastTrackForImportsPhase(firDeclarationToResolve, moduleFileCache, checkPCE)) {
-                return firDeclarationToResolve
+                return
            }
        }
        when (firDeclarationToResolve) {
@@ -249,9 +245,8 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                    scopeSession,
                    toPhase,
                    checkPCE,
-                    skipLocalDeclaration
                )
-                return firDeclarationToResolve
+                return
            }

            is FirBackingField -> {
@@ -261,7 +256,6 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                    scopeSession,
                    toPhase,
                    checkPCE,
-                    skipLocalDeclaration
                )
            }

@@ -273,51 +267,61 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                    scopeSession = scopeSession,
                    checkPCE = checkPCE,
                )
-                return firDeclarationToResolve
+                return
            }
+            else -> {}
        }

-        if (!firDeclarationToResolve.isValidForResolve()) return firDeclarationToResolve
-        if (firDeclarationToResolve.resolvePhase >= toPhase) return firDeclarationToResolve
+        if (!firDeclarationToResolve.isValidForResolve()) return
+        if (firDeclarationToResolve.resolvePhase >= toPhase) return


        val requestedDeclarationDesignation = firDeclarationToResolve.tryCollectDesignationWithFile()

        val designation: FirDeclarationDesignationWithFile
        val neededPhase: FirResolvePhase
-        val isLocalDeclarationResolveRequested: Boolean
+
        if (requestedDeclarationDesignation != null) {
            designation = requestedDeclarationDesignation
            neededPhase = toPhase
-            isLocalDeclarationResolveRequested = false
        } else {
            val possiblyLocalDeclaration = firDeclarationToResolve.getKtDeclarationForFirElement()
            val nonLocalDeclaration = possiblyLocalDeclaration.getNonLocalContainingOrThisDeclaration()
                ?: error("Container for local declaration cannot be null")
-
-            isLocalDeclarationResolveRequested =
+            val isLocalDeclarationResolveRequested =
                possiblyLocalDeclaration != nonLocalDeclaration
-            if (isLocalDeclarationResolveRequested && skipLocalDeclaration) return firDeclarationToResolve

-            val nonLocalFirDeclaration = nonLocalDeclaration.findSourceNonLocalFirDeclaration(
-                firFileBuilder,
-                firDeclarationToResolve.moduleData.session.firIdeProvider.symbolProvider,
-                moduleFileCache
-            )
+            val declarationToResolve: FirDeclaration

-            neededPhase = if (isLocalDeclarationResolveRequested) FirResolvePhase.BODY_RESOLVE else toPhase
+            if (isLocalDeclarationResolveRequested) {
+                val enumEntry = possiblyLocalDeclaration.getContainingEnumEntryAsMemberOfEnumEntry() ?: return

-            if (nonLocalFirDeclaration.resolvePhase >= neededPhase) return firDeclarationToResolve
-            if (!nonLocalFirDeclaration.isValidForResolve()) return firDeclarationToResolve
+                declarationToResolve = enumEntry.findSourceNonLocalFirDeclaration(
+                    firFileBuilder,
+                    firDeclarationToResolve.moduleData.session.firIdeProvider.symbolProvider,
+                    moduleFileCache
+                )
+                neededPhase = FirResolvePhase.BODY_RESOLVE
+            } else {
+                declarationToResolve = nonLocalDeclaration.findSourceNonLocalFirDeclaration(
+                    firFileBuilder,
+                    firDeclarationToResolve.moduleData.session.firIdeProvider.symbolProvider,
+                    moduleFileCache
+                )
+                neededPhase = toPhase
+            }

-            designation = nonLocalFirDeclaration.collectDesignationWithFile()
+            if (declarationToResolve.resolvePhase >= neededPhase) return
+            if (!declarationToResolve.isValidForResolve()) return
+
+            designation = declarationToResolve.collectDesignationWithFile()
        }

-        if (designation.declaration.resolvePhase >= neededPhase) return firDeclarationToResolve
+        if (designation.declaration.resolvePhase >= neededPhase) return

        if (neededPhase == FirResolvePhase.IMPORTS) {
            resolveFileToImports(designation.firFile, moduleFileCache, checkPCE)
-            return firDeclarationToResolve
+            return
        }

        moduleFileCache.firFileLockProvider.runCustomResolveUnderLock(designation.firFile, checkPCE) {
@@ -332,33 +336,8 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
                designation.declaration
            }
        }
-
-        if (!isLocalDeclarationResolveRequested) return firDeclarationToResolve
-        return remapDeclarationInContainerIfNeeded(firDeclarationToResolve, designation.declaration, toPhase)
    }

-    private fun <D : FirDeclaration> remapDeclarationInContainerIfNeeded(
-        declarationToRemap: D,
-        firContainer: FirDeclaration,
-        firResolvePhase: FirResolvePhase
-    ): D {
-        if (declarationToRemap.resolvePhase >= firResolvePhase) return declarationToRemap
-        val realPsi = declarationToRemap.realPsi
-        check(realPsi != null) {
-            "Cannot remap element without PSI"
-        }
-        @Suppress("UNCHECKED_CAST")
-        val firDeclaration = FirElementFinder.findElementIn<FirDeclaration>(firContainer) {
-            it.realPsi == realPsi
-        } as? D
-        check(firDeclaration != null) {
-            "Containing declaration was resolved but local didn't found in it"
-        }
-        check(firDeclaration.resolvePhase >= firResolvePhase) {
-            "Found local declaration wasn't completely resolved"
-        }
-        return firDeclaration
-    }

    private fun runLazyDesignatedResolveWithoutLock(
        designation: FirDeclarationDesignationWithFile,
@@ -433,10 +412,8 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
    }
 }

-private fun KtDeclaration.isMemberOfEnumEntry(container: KtDeclaration) : Boolean {
-    if (container !is KtClass) return false
-    if (!container.isEnum()) return false
-    val enumEntryBody = this.parent as? KtClassBody ?: return false
-    return enumEntryBody.parent == container
+private fun KtDeclaration.getContainingEnumEntryAsMemberOfEnumEntry(): KtEnumEntry? {
+    val body = parent as? KtClassBody ?: return null
+    return body.parent as? KtEnumEntry
 }