Dirty attempt of custom on-air designation
diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirScopeProvider.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirScopeProvider.kt
index 6bc12dd..938c2f4 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirScopeProvider.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirScopeProvider.kt
@@ -256,17 +256,16 @@
originalFile: KtFile,
positionInFakeFile: KtElement
): KtScopeContext {
- val fakeFile = positionInFakeFile.containingKtFile
-
// If the position is in KDoc, we want to pass the owning declaration to the ContextCollector.
// That way, the resulting scope will contain all the nested declarations which can be references by KDoc.
val parentKDoc = positionInFakeFile.parentOfType<KDoc>()
val correctedPosition = parentKDoc?.owner ?: positionInFakeFile
- val context = ContextCollector.process(
- fakeFile.getOrBuildFirFile(firResolveSession),
+ val context = ContextCollector.processFake(
+ originalFile.getOrBuildFirFile(firResolveSession),
SessionHolderImpl(analysisSession.useSiteSession, getScopeSession()),
correctedPosition,
+ firResolveSession,
)
val towerDataContext =
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
index 07cf3ca..2709efa 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
@@ -278,6 +278,15 @@
return fileAnnotationsContainer.annotations.single()
}
+ fun getDeclarationForOnAirResolve(originalPlace: PsiElement): KtElement? {
+ val minimalOriginalDeclarationToReplace = originalPlace.getNonLocalContainingOrThisDeclarationCodeFragmentAware {
+ it.isApplicableForOnAirResolve()
+ }
+
+ val originalDeclaration = minimalOriginalDeclarationToReplace?.onAirGetNonLocalContainingOrThisDeclaration()
+ return originalDeclaration
+ }
+
/**
* Creates a new fir declaration from closer non-local declaration based on [originalPlace] position.
* The resulted [FirResolvePhase] depends on [forcedResolvePhase] or the position of [originalPlace] inside the non-local declaration.
@@ -355,7 +364,7 @@
session.moduleComponents.firModuleLazyDeclarationResolver.runLazyDesignatedOnAirResolve(
FirDesignationWithFile(originalDesignationPath.path, newFirDeclaration, originalFirFile),
collector,
- forcedResolvePhase ?: requiredResolvePhase(minimalOriginalDeclarationToReplace, originalPlace),
+ resolvePhase = forcedResolvePhase ?: requiredResolvePhase(minimalOriginalDeclarationToReplace, originalPlace),
)
return newFirDeclaration
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt
index dfef535..4baeb17 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt
@@ -7,8 +7,12 @@
import com.intellij.openapi.progress.ProgressManager
import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithScript
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirSafe
import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.ContextKind
import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.Context
@@ -36,6 +40,8 @@
import org.jetbrains.kotlin.fir.types.typeContext
import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
import org.jetbrains.kotlin.util.PrivateForInline
import org.jetbrains.kotlin.utils.yieldIfNotNull
@@ -106,6 +112,53 @@
return null
}
+ fun processFake(file: FirFile, holder: SessionHolder, targetElement: PsiElement, resolveSession: LLFirResolveSession, bodyElement: PsiElement? = targetElement): Context? {
+ val isBodyContextCollected = bodyElement != null
+ val acceptedElements = targetElement.parentsWithSelf.flatMap {
+ val originalDeclarations = when (it) {
+ is KtDeclaration -> it.originalDeclaration
+ is KtFile -> it.originalFile
+ else -> null
+ }
+ listOfNotNull(it, originalDeclarations)
+ }.toSet()
+
+ val contextProvider = process(file, holder, computeFakeDesignation(file, targetElement, resolveSession), isBodyContextCollected) { candidate ->
+ when (candidate) {
+ targetElement -> FilterResponse.STOP
+ in acceptedElements -> FilterResponse.CONTINUE
+ else -> FilterResponse.SKIP
+ }
+ }
+
+ for (acceptedElement in acceptedElements) {
+ if (acceptedElement === bodyElement) {
+ val bodyContext = contextProvider[acceptedElement, ContextKind.BODY]
+ if (bodyContext != null) {
+ return bodyContext
+ }
+ }
+
+ val elementContext = contextProvider[acceptedElement, ContextKind.SELF]
+ if (elementContext != null) {
+ return elementContext
+ }
+
+ val original = when (acceptedElement) {
+ is KtDeclaration -> acceptedElement.originalDeclaration
+ is KtFile -> acceptedElement.originalFile
+ else -> null
+ } ?: continue
+
+ val originalElementContext = contextProvider[original, ContextKind.SELF]
+ if (originalElementContext != null) {
+ return originalElementContext
+ }
+ }
+
+ return null
+ }
+
fun computeDesignation(file: FirFile, targetElement: PsiElement): FirDesignation? {
val contextKtDeclaration = targetElement.getNonLocalContainingOrThisDeclaration()
if (contextKtDeclaration != null) {
@@ -124,6 +177,26 @@
return null
}
+ fun computeFakeDesignation(file: FirFile, targetElement: PsiElement, resolveSession: LLFirResolveSession): FirDesignation? {
+ val fakeContextKtDeclaration = LowLevelFirApiFacadeForResolveOnAir.getDeclarationForOnAirResolve(targetElement) as? KtDeclaration
+ val fakeFirDeclaration = fakeContextKtDeclaration?.getOrBuildFirSafe<FirDeclaration>(resolveSession) ?: return null
+ val contextKtDeclaration = PsiTreeUtil.findSameElementInCopy(fakeContextKtDeclaration, file.psi as KtFile)
+ if (contextKtDeclaration != null) {
+ val designationPath = FirElementFinder.collectDesignationPath(file, contextKtDeclaration)
+ if (designationPath != null) {
+ val script = file.declarations.singleOrNull() as? FirScript
+
+ return if (script == null || script === designationPath.target) {
+ FirDesignation(designationPath.path, fakeFirDeclaration)
+ } else {
+ FirDesignationWithScript(designationPath.path, designationPath.target, script)
+ }
+ }
+ }
+
+ return null
+ }
+
/**
* Processes the [FirFile], collecting contexts for elements matching the [filter].
*
