diff --git a/analysis/low-level-api-fir/build.gradle.kts b/analysis/low-level-api-fir/build.gradle.kts index 7ab6219912a..9e8cec824c1 100644 --- a/analysis/low-level-api-fir/build.gradle.kts +++ b/analysis/low-level-api-fir/build.gradle.kts @@ -76,6 +76,7 @@ allprojects { kotlinOptions { freeCompilerArgs += "-opt-in=org.jetbrains.kotlin.fir.symbols.SymbolInternals" freeCompilerArgs += "-opt-in=org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals" + freeCompilerArgs += "-opt-in=org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals" } } } diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirNonLocalDeclarationSearcher.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirNonLocalDeclarationSearcher.kt index d640d31c0f8..49a368b48da 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirNonLocalDeclarationSearcher.kt +++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirNonLocalDeclarationSearcher.kt @@ -10,80 +10,104 @@ import org.jetbrains.kotlin.fir.FirElement import org.jetbrains.kotlin.fir.declarations.* import org.jetbrains.kotlin.fir.declarations.utils.classId import org.jetbrains.kotlin.fir.packageFqName +import org.jetbrains.kotlin.fir.psi import org.jetbrains.kotlin.fir.render import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid import org.jetbrains.kotlin.name.ClassId +import org.jetbrains.kotlin.psi.KtDeclaration +import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject object FirElementFinder { fun findClassifierWithClassId( firFile: FirFile, classId: ClassId, - ): FirClassLikeDeclaration? = findPathToClassifierWithClassId( + ): FirClassLikeDeclaration? = collectDesignationPath( firFile = firFile, - classId = classId, - expectedMemberDeclaration = null - )?.second + containerClassId = classId.outerClassId, + expectedDeclarationAcceptor = { it is FirClassLikeDeclaration && it.symbol.name == classId.shortClassName }, + )?.target?.let { it as FirClassLikeDeclaration } fun findClassPathToDeclaration( firFile: FirFile, declarationContainerClassId: ClassId, targetMemberDeclaration: FirDeclaration, - ): List? = findPathToClassifierWithClassId( + ): List? = collectDesignationPath( firFile = firFile, - classId = declarationContainerClassId, - expectedMemberDeclaration = targetMemberDeclaration - )?.first + containerClassId = declarationContainerClassId, + expectedDeclarationAcceptor = { it == targetMemberDeclaration }, + )?.path - private fun findPathToClassifierWithClassId( + fun findDeclaration(firFile: FirFile, nonLocalDeclaration: KtDeclaration): FirDeclaration? = collectDesignationPath( + firFile = firFile, + nonLocalDeclaration = nonLocalDeclaration, + )?.target + + fun findPathToDeclarationWithTarget( firFile: FirFile, - classId: ClassId, - expectedMemberDeclaration: FirDeclaration?, - ): Pair, FirClassLikeDeclaration>? { - requireWithAttachmentBuilder(!classId.isLocal, { "ClassId should not be local" }) { - withEntry("classId", classId) { it.asString() } - } - requireWithAttachmentBuilder( - firFile.packageFqName == classId.packageFqName, - { "ClassId should not be local" } - ) { - withEntry("FirFile.packageName", firFile.packageFqName) { it.asString() } - withEntry("ClassId.packageName", classId.packageFqName) { it.asString() } + nonLocalDeclaration: KtDeclaration, + ): List? = collectDesignationPath( + firFile = firFile, + nonLocalDeclaration = nonLocalDeclaration, + )?.pathWithTarget + + private class FirDeclarationDesignation( + val path: List, + val target: FirDeclaration, + ) { + val pathWithTarget: List get() = path + target + } + + private fun collectDesignationPath( + firFile: FirFile, + nonLocalDeclaration: KtDeclaration, + ): FirDeclarationDesignation? = collectDesignationPath( + firFile = firFile, + containerClassId = nonLocalDeclaration.containingClassOrObject?.getClassId(), + expectedDeclarationAcceptor = { it.psi == nonLocalDeclaration }, + ) + + private fun collectDesignationPath( + firFile: FirFile, + containerClassId: ClassId?, + expectedDeclarationAcceptor: (FirDeclaration) -> Boolean, + ): FirDeclarationDesignation? { + if (containerClassId != null) { + requireWithAttachmentBuilder(!containerClassId.isLocal, { "ClassId should not be local" }) { + withEntry("classId", containerClassId) { it.asString() } + } + + requireWithAttachmentBuilder( + firFile.packageFqName == containerClassId.packageFqName, + { "ClassId should not be local" } + ) { + withEntry("FirFile.packageName", firFile.packageFqName) { it.asString() } + withEntry("ClassId.packageName", containerClassId.packageFqName) { it.asString() } + } } - val classIdPathSegment = classId.relativeClassName.pathSegments() + val classIdPathSegment = containerClassId?.relativeClassName?.pathSegments().orEmpty() val path = ArrayList(classIdPathSegment.size) - var result: FirClassLikeDeclaration? = null + var result: FirDeclaration? = null fun find(declarations: Iterable, classIdPathIndex: Int): Boolean { - val currentClassSegment = classIdPathSegment[classIdPathIndex] - + val currentClassSegment = classIdPathSegment.getOrNull(classIdPathIndex) for (subDeclaration in declarations) { - if (subDeclaration is FirScript) { - val scriptDeclarations = subDeclaration.statements.asSequence().filterIsInstance() - if (find(scriptDeclarations.asIterable(), classIdPathIndex)) { + when { + currentClassSegment == null && expectedDeclarationAcceptor(subDeclaration) -> { + result = subDeclaration return true } - continue - } - - if (subDeclaration is FirClassLikeDeclaration && currentClassSegment == subDeclaration.symbol.name) { - if (classIdPathIndex == classIdPathSegment.lastIndex) { - if (expectedMemberDeclaration == null || - subDeclaration is FirRegularClass && expectedMemberDeclaration in subDeclaration.declarations - ) { - if (subDeclaration is FirRegularClass) { - path += subDeclaration - } - - result = subDeclaration + subDeclaration is FirScript -> { + val scriptDeclarations = subDeclaration.statements.asSequence().filterIsInstance() + if (find(scriptDeclarations.asIterable(), classIdPathIndex)) { return true } continue } - if (subDeclaration is FirRegularClass) { + subDeclaration is FirRegularClass && currentClassSegment == subDeclaration.symbol.name -> { path += subDeclaration if (find(subDeclaration.declarations, classIdPathIndex + 1)) { return true @@ -99,7 +123,10 @@ object FirElementFinder { find(firFile.declarations, classIdPathIndex = 0) return result?.let { - (path as List) to it + FirDeclarationDesignation( + path = if (path.isEmpty()) emptyList() else path, + target = it, + ) } } diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt index c82e0f6429e..acf965d3fed 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt +++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt @@ -5,7 +5,10 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.util -import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError +import com.intellij.psi.PsiElement +import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals +import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession +import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.canBePartOfParentDeclaration import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.containingDeclaration import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder @@ -29,39 +32,59 @@ internal fun KtDeclaration.findSourceNonLocalFirDeclaration( provider: FirProvider, containerFirFile: FirFile? = null, ): FirDeclaration { - //TODO test what way faster - findSourceNonLocalFirDeclarationByProvider(firFileBuilder, provider, containerFirFile)?.let { return it } - findSourceByTraversingWholeTree(firFileBuilder, containerFirFile)?.let { return it } + val firFile = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile) + + // TODO test what way faster + if (isPhysical) { + // do not request providers with non-physical psi in order not to leak them there and + // to avoid inconsistency between physical psi and its copy during completion + findSourceNonLocalFirDeclarationByProvider( + firDeclarationProvider = { declaration -> + if (declaration is KtClassLikeDeclaration) { + declaration.findFir(provider) + } else { + val containingClassOrObject = declaration.containingClassOrObject + val declarations = if (containingClassOrObject != null) { + val containerClassFir = containingClassOrObject.findFir(provider) as? FirRegularClass + containerClassFir?.declarations + } else { + if (declaration.containingKtFile.isScript()) { + // .kts will have a single [FirScript] as a declaration. We need to unwrap statements in it. + (firFile.declarations.singleOrNull() as? FirScript)?.statements?.filterIsInstance() + } else { + firFile.declarations + } + } + + // It is possible that we will not be able to find the needed declaration here when the code is invalid + // e.g., we have two conflicting declarations with the same name, + // and we are searching for the wrong one + declarations?.find { it.psi == declaration } + } + }, + )?.let { return it } + } + + findSourceNonLocalFirDeclarationByProvider( + firDeclarationProvider = { declaration -> + FirElementFinder.findDeclaration(firFile, declaration) + }, + )?.let { return it } + errorWithFirSpecificEntries( "No fir element was found for", psi = this, - fir = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile), + fir = firFile, additionalInfos = { withEntry("isPhysical", isPhysical.toString()) } ) } -internal fun KtDeclaration.findFirDeclarationForAnyFirSourceDeclaration( - firFileBuilder: LLFirFileBuilder, - provider: FirProvider, -): FirDeclaration { - val nonLocalDeclaration = getNonLocalContainingOrThisDeclaration() - ?.findSourceNonLocalFirDeclaration(firFileBuilder, provider) - ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile) - val originalDeclaration = originalDeclaration - val fir = FirElementFinder.findElementIn(nonLocalDeclaration) { firDeclaration -> - firDeclaration.psi == this || firDeclaration.psi == originalDeclaration - } - return fir - ?: errorWithFirSpecificEntries("FirDeclaration was not found", psi = this) -} - -internal inline fun KtDeclaration.findFirDeclarationForAnyFirSourceDeclarationOfType( - firFileBuilder: LLFirFileBuilder, - provider: FirProvider, -): FirDeclaration { - val fir = findFirDeclarationForAnyFirSourceDeclaration(firFileBuilder, provider) - if (fir !is F) throwUnexpectedFirElementError(fir, this, F::class) - return fir +@KtAnalysisApiInternals +fun PsiElement.collectContainingDeclarationsIfNonLocal(session: LLFirResolveSession): List? { + val ktFile = containingFile as? KtFile ?: return null + val ktDeclaration = getNonLocalContainingOrThisDeclaration { !it.canBePartOfParentDeclaration } ?: return null + val firFile = session.getOrBuildFirFile(ktFile) + return FirElementFinder.findPathToDeclarationWithTarget(firFile, ktDeclaration) } internal fun KtElement.findSourceByTraversingWholeTree( @@ -81,56 +104,22 @@ internal fun KtElement.findSourceByTraversingWholeTree( } private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider( - firFileBuilder: LLFirFileBuilder, - provider: FirProvider, - containerFirFile: FirFile?, + firDeclarationProvider: (KtDeclaration) -> FirDeclaration?, ): FirDeclaration? { - if (!isPhysical) { - //do not request providers with non-physical psi in order not to leak them there and - //to avoid inconsistency between physical psi and its copy during completion - return null - } - val candidate = when { - this is KtClassOrObject -> findFir(provider) - this is KtNamedDeclaration && (this is KtProperty || this is KtNamedFunction) -> { - val containerClass = containingClassOrObject - val declarations = if (containerClass != null) { - val containerClassFir = containerClass.findFir(provider) as? FirRegularClass - containerClassFir?.declarations - } else { - val ktFile = containingKtFile - val firFile = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(ktFile) - if (ktFile.isScript()) { - // .kts will have a single [FirScript] as a declaration. We need to unwrap statements in it. - (firFile.declarations.singleOrNull() as? FirScript)?.statements?.filterIsInstance() - } else { - firFile.declarations - } - } + val candidate = when (this) { + is KtClassOrObject, + is KtProperty, + is KtNamedFunction, + is KtConstructor<*>, + is KtClassInitializer, + is KtTypeAlias, + is KtDestructuringDeclaration, + is KtScript, + -> firDeclarationProvider(this) - /* - It is possible that we will not be able to find needed declaration here when the code is invalid, - e.g, we have two conflicting declarations with the same name and we are searching in the wrong one - */ - declarations?.firstOrNull { it.psi == this } - } - this is KtConstructor<*> || this is KtClassInitializer -> { - val containingClass = containingClassOrObject - ?: errorWithFirSpecificEntries("Container class should be not null for KtConstructor", psi = this) - val containerClassFir = containingClass.findFir(provider) as? FirRegularClass ?: return null - containerClassFir.declarations.firstOrNull { it.psi == this } - } - this is KtTypeAlias -> findFir(provider) - this is KtDestructuringDeclaration -> { - val firFile = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile) - firFile.declarations.firstOrNull { it.psi == this } - } - this is KtScript -> containerFirFile?.declarations?.singleOrNull { it is FirScript } - this is KtPropertyAccessor -> { + is KtPropertyAccessor -> { val firPropertyDeclaration = property.findSourceNonLocalFirDeclarationByProvider( - firFileBuilder, - provider, - containerFirFile, + firDeclarationProvider, ) as? FirVariable ?: return null if (isGetter) { @@ -139,32 +128,36 @@ private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider( firPropertyDeclaration.setter } } - this is KtParameter -> { - val ownerFunction = ownerFunction - ?: errorWithFirSpecificEntries("Containing function should be not null for KtParameter", psi = this) + + is KtParameter -> { + val ownerFunction = ownerFunction ?: errorWithFirSpecificEntries( + "Containing function should be not null for KtParameter", + psi = this, + ) val firFunctionDeclaration = ownerFunction.findSourceNonLocalFirDeclarationByProvider( - firFileBuilder, - provider, - containerFirFile, + firDeclarationProvider, ) as? FirFunction ?: return null firFunctionDeclaration.valueParameters[parameterIndex()] } - this is KtTypeParameter -> { - val declaration = containingDeclaration - ?: errorWithFirSpecificEntries("Containing declaration should be not null for KtTypeParameter", psi = this) + + is KtTypeParameter -> { + val declaration = containingDeclaration ?: errorWithFirSpecificEntries( + "Containing declaration should be not null for KtTypeParameter", + psi = this, + ) val firTypeParameterOwner = declaration.findSourceNonLocalFirDeclarationByProvider( - firFileBuilder, - provider, - containerFirFile, + firDeclarationProvider, ) as? FirTypeParameterRefsOwner ?: return null firTypeParameterOwner.typeParameters.firstOrNull { it.psi == this } as FirDeclaration } + else -> errorWithFirSpecificEntries("Invalid container", psi = this) } + return candidate?.takeIf { it.realPsi == this } }