[Analysis API] add ability for a custom element navigation for a KtResolveExtensionFile

^KT-57930

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirReference.kt

@@ -13,11 +13,15 @@ import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.psiTargetsProvider
+import org.jetbrains.kotlin.analysis.project.structure.KtModuleStructureInternals
+import org.jetbrains.kotlin.psi.KtFile
 
 interface KtFirReference : KtReference, KtSymbolBasedReference {
     fun getResolvedToPsi(analysisSession: KtAnalysisSession, referenceTargetSymbols: Collection<KtSymbol>): Collection<PsiElement> =
         with(analysisSession) {
             referenceTargetSymbols.flatMap { symbol ->
+                symbol.getPsiForGeneratedExtensionSymbol()?.let { return@flatMap it }
                 when (symbol) {
                     is KtFirSymbol<*> -> getPsiDeclarations(symbol)
                     else -> listOfNotNull(symbol.psi)
@@ -25,6 +29,19 @@ interface KtFirReference : KtReference, KtSymbolBasedReference {
             }
         }
 
+    context(KtAnalysisSession)
+    @OptIn(KtModuleStructureInternals::class)
+    private fun KtSymbol.getPsiForGeneratedExtensionSymbol(): Collection<PsiElement>? {
+        // TODO symbols generated by KtResolveExtension should probably have their own KtSymbolOrigin, now it's KtSymbolOrigin.SOURCE
+        if (origin != KtSymbolOrigin.SOURCE) return null
+        val psi = psi ?: return null
+        val ktFile = psi.containingFile as? KtFile ?: return null
+        val psiTargetsProvider = ktFile.virtualFile?.psiTargetsProvider ?: return null
+        return with(psiTargetsProvider) {
+            getReferenceTargetsForSymbol(this@getPsiForGeneratedExtensionSymbol)
+        }
+    }
+
     fun getResolvedToPsi(analysisSession: KtAnalysisSession): Collection<PsiElement> =
         with(analysisSession) {
             getResolvedToPsi(analysisSession, resolveToSymbols())

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/util/KtResolveExtensionForTest.kt

@@ -7,9 +7,13 @@ package org.jetbrains.kotlin.analysis.api.impl.base.test.util
 
 import com.intellij.mock.MockProject
 import com.intellij.openapi.util.ModificationTracker
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtension
 import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionFile
 import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionProvider
+import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionReferencePsiTargetsProvider
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.test.framework.services.environmentManager
 import org.jetbrains.kotlin.name.FqName
@@ -66,4 +70,12 @@ class KtResolveExtensionFileForTests(
     override fun getTopLevelCallableNames(): Set<Name> = topLevelCallableNames
 
     override fun buildFileText(): String = fileText
+
+    override fun createPsiTargetsProvider(): KtResolveExtensionReferencePsiTargetsProvider {
+        return object : KtResolveExtensionReferencePsiTargetsProvider() {
+            override fun KtAnalysisSession.getReferenceTargetsForSymbol(symbol: KtSymbol): Collection<PsiElement> {
+                return emptyList()
+            }
+        }
+    }
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/resolve/extensions/KtResolveExtensionFile.kt

@@ -77,11 +77,11 @@ abstract class KtResolveExtensionFile {
     abstract fun buildFileText(): String
 
     /**
-     * Creates a [KtResolveExtensionReferenceTargetsPsiProvider] for this [KtResolveExtensionFile].
+     * Creates a [KtResolveExtensionReferencePsiTargetsProvider] for this [KtResolveExtensionFile].
      *
-     * @see KtResolveExtensionReferenceTargetsPsiProvider
+     * @see KtResolveExtensionReferencePsiTargetsProvider
      * @see KtResolveExtensionFile
      */
-    abstract fun createTargetPsiProvider(): KtResolveExtensionReferenceTargetsPsiProvider
+    abstract fun createPsiTargetsProvider(): KtResolveExtensionReferencePsiTargetsProvider
 
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/resolve/extensions/KtResolveExtensionReferencePsiTargetsProvider.kt

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.api.resolve.extensions
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
+
+abstract class KtResolveExtensionReferencePsiTargetsProvider {
+    /**
+     * Provides a [PsiElement] where `reference.resolveTo` will lead for a [symbol]
+     *
+     * Usually returns a single result. Might return an empty collection if there is no navigation target.
+     * Also, might multiple targets in a case of ambiguity or multiple targets for a [symbol]
+     *
+     * Returned [PsiElement] will be used as a navigation target for a reference inside the IDE.
+     */
+    abstract fun KtAnalysisSession.getReferenceTargetsForSymbol(symbol: KtSymbol): Collection<PsiElement>
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolve/extensions/LLFirResolveExtensionTool.kt

@@ -5,10 +5,12 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions
 
+import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.ModificationTracker
+import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtension
 import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionFile
-import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionReferenceTargetsPsiProvider
+import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionReferencePsiTargetsProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.FileBasedKotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirSymbolProviderNameCache
@@ -192,16 +194,24 @@ private class LLFirResolveExtensionToolDeclarationProvider(
                 eventSystemEnabled = true // so every generated KtFile backed by some VirtualFile
             )
             val text = file.buildFileText()
-            val ktFile = createKtFile(factory, file.getFileName(), text)
+            val psiTargetsProvider = file.createPsiTargetsProvider()
+            val ktFile = createKtFile(factory, file.getFileName(), text, psiTargetsProvider)
             FileBasedKotlinDeclarationProvider(ktFile)
         }
     }
 
 
     @OptIn(KtModuleStructureInternals::class)
-    private fun createKtFile(factory: KtPsiFactory, fileName: String, fileText: String): KtFile {
+    private fun createKtFile(
+        factory: KtPsiFactory,
+        fileName: String,
+        fileText: String,
+        psiTargetsProvider: KtResolveExtensionReferencePsiTargetsProvider
+    ): KtFile {
         val ktFile = factory.createFile(fileName, fileText)
-        ktFile.virtualFile.analysisExtensionFileContextModule = ktModule
+        val virtualFile = ktFile.virtualFile
+        virtualFile.analysisExtensionFileContextModule = ktModule
+        virtualFile.psiTargetsProvider = psiTargetsProvider
         return ktFile
     }
 }
@@ -274,3 +284,6 @@ private fun KtResolveExtensionFile.mayHaveTopLevelClassifier(name: Name): Boolea
 private fun KtResolveExtensionFile.mayHaveTopLevelCallable(name: Name): Boolean {
     return name in getTopLevelCallableNames()
 }
+
+@KtModuleStructureInternals
+public var VirtualFile.psiTargetsProvider: KtResolveExtensionReferencePsiTargetsProvider? by UserDataProperty(Key.create("KT_RESOLVE_EXTENSION_PSI_TARGETS_PROVIDER"))