FIR IDE: introduce SOURCE_MEMBER_GENERATED declaration kind

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtSymbol.kt

@@ -17,5 +17,14 @@ interface KtSymbol : ValidityTokenOwner {
 }
 
 enum class KtSymbolOrigin {
-    SOURCE, LIBRARY, JAVA, SAM_CONSTRUCTOR
+    SOURCE,
+
+    /**
+     * Declaration which do not have it's PSI source and was generated, they are:
+     * For data classes the `copy`, `component{N}`, `toString`, `equals`, `hashCode` functions are generated
+     * For enum classes the `valueOf` & `values` functions are generated
+     */
+    SOURCE_MEMBER_GENERATED,
+    LIBRARY,
+    JAVA, SAM_CONSTRUCTOR
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirIdeDeserializedDeclarationSourceProvider.kt

@@ -10,6 +10,8 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSession
 import org.jetbrains.kotlin.idea.stubindex.KotlinFullClassNameIndex
 import org.jetbrains.kotlin.idea.stubindex.KotlinTopLevelFunctionFqnNameIndex
@@ -138,7 +140,15 @@ object FirIdeDeserializedDeclarationSourceProvider {
 private fun KtElement.isCompiled(): Boolean = containingKtFile.isCompiled
 
 fun FirElement.findPsi(project: Project): PsiElement? =
-    psi ?: FirIdeDeserializedDeclarationSourceProvider.findPsi(this, project)
+    realPsi ?: FirIdeDeserializedDeclarationSourceProvider.findPsi(this, project)
 
 fun FirElement.findPsi(session: FirSession): PsiElement? =
     findPsi((session as FirIdeSession).project)
+
+/**
+ * Finds [PsiElement] which will be used as go-to referenced element for [KtPsiReference]
+ * For data classes & enums generated members like `copy` `componentN`, `values` it will return corresponding enum/data class
+ * Otherwise, behaves the same way as [findPsi] returns exact PSI declaration corresponding to passed [FirDeclaration]
+ */
+fun FirDeclaration.findReferencePsi(): PsiElement? =
+    psi ?: FirIdeDeserializedDeclarationSourceProvider.findPsi(this, (session as FirIdeSession).project)

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirSymbolContainingDeclarationProvider.kt

@@ -23,7 +23,8 @@ internal class KtFirSymbolContainingDeclarationProvider(
         if (symbol is KtPackageSymbol) return null
         if (symbol.symbolKind == KtSymbolKind.TOP_LEVEL) return null
         return when (symbol.origin) {
-            KtSymbolOrigin.SOURCE -> getContainingDeclarationForKotlinInSourceSymbol(symbol)
+            KtSymbolOrigin.SOURCE, KtSymbolOrigin.SOURCE_MEMBER_GENERATED ->
+                getContainingDeclarationForKotlinInSourceSymbol(symbol)
             KtSymbolOrigin.LIBRARY -> getContainingDeclarationForLibrarySymbol(symbol)
             KtSymbolOrigin.JAVA -> TODO()
             KtSymbolOrigin.SAM_CONSTRUCTOR -> TODO()
@@ -31,11 +32,17 @@ internal class KtFirSymbolContainingDeclarationProvider(
     }
 
     private fun getContainingDeclarationForKotlinInSourceSymbol(symbol: KtSymbolWithKind): KtSymbolWithKind = with(analysisSession) {
-        require(symbol.origin == KtSymbolOrigin.SOURCE)
+        require(symbol.origin == KtSymbolOrigin.SOURCE || symbol.origin == KtSymbolOrigin.SOURCE_MEMBER_GENERATED)
         val psi = symbol.psi ?: error("PSI should present for declaration built by Kotlin code")
         check(psi is KtDeclaration) { "PSI of kotlin declaration should be KtDeclaration" }
-        val containingDeclaration = psi.parentOfType<KtDeclaration>()
+        val containingDeclaration = when (symbol.origin) {
-            ?: error("Containing declaration should present for non-toplevel declaration")
+            KtSymbolOrigin.SOURCE -> psi.parentOfType()
+                ?: error("Containing declaration should present for non-toplevel declaration")
+            KtSymbolOrigin.SOURCE_MEMBER_GENERATED -> psi
+            else -> error("Unsupported declaration origin ${symbol.origin}")
+        }
+
+
         return with(analysisSession) {
             val containingSymbol = containingDeclaration.getSymbol()
             check(containingSymbol is KtSymbolWithKind)

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbol.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
+import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
@@ -22,7 +23,13 @@ internal interface KtFirSymbol<F : FirDeclaration> : KtSymbol, ValidityTokenOwne
 
 
 private tailrec fun FirDeclaration.ktSymbolOrigin(): KtSymbolOrigin = when (origin) {
-    FirDeclarationOrigin.Source -> KtSymbolOrigin.SOURCE
+    FirDeclarationOrigin.Source -> {
+        if (source?.kind == FirFakeSourceElementKind.DataClassGeneratedMembers
+            || source?.kind == FirFakeSourceElementKind.EnumGeneratedDeclaration
+        ) {
+            KtSymbolOrigin.SOURCE_MEMBER_GENERATED
+        } else KtSymbolOrigin.SOURCE
+    }
     FirDeclarationOrigin.Library -> KtSymbolOrigin.LIBRARY
     FirDeclarationOrigin.Java -> KtSymbolOrigin.JAVA
     FirDeclarationOrigin.SamConstructor -> KtSymbolOrigin.SAM_CONSTRUCTOR

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReference.kt

@@ -6,13 +6,18 @@
 package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.fir.findReferencePsi
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtSymbolBasedReference
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 
 interface KtFirReference : KtReference, KtSymbolBasedReference {
     fun getResolvedToPsi(analysisSession: KtAnalysisSession): Collection<PsiElement> =
-        analysisSession.resolveToSymbols().mapNotNull(KtSymbol::psi)
+        analysisSession.resolveToSymbols().mapNotNull { symbol ->
+            (symbol as? KtFirSymbol<*>)?.firRef?.withFir { it.findReferencePsi() }
+                ?: symbol.psi
+        }
 
     override val resolver get() = KtFirReferenceResolver
 }