[LL] ContextCollector: dump context only for real source elements

Some fake sources might be traversed and cached before the real element,
so an incorrect result might be returned.
That was the case for enum classes. They have a synthetic super type
reference with the fake (`EnumSuperTypeRef`) source which contains
the class psi. As a result, the incorrect context was captured and
returned for the class

^KT-79653 Fixed

Author: dimonchik0036

analysis/analysis-api/testData/components/referenceShortener/shortenRange/nestedClasses/classHeaderPositions/constructorParameterVsTopLevelProperty_noConflict.txt

@@ -14,16 +14,20 @@ with default settings:
 [qualifier] test.bigProperty
 [qualifier] test.bigProperty
 [qualifier] test.bigProperty
+[qualifier] test.bigProperty
 with DO_NOT_SHORTEN:
 with SHORTEN_IF_ALREADY_IMPORTED:
 [qualifier] test.bigProperty
 [qualifier] test.bigProperty
 [qualifier] test.bigProperty
+[qualifier] test.bigProperty
 with SHORTEN_AND_IMPORT:
 [qualifier] test.bigProperty
 [qualifier] test.bigProperty
 [qualifier] test.bigProperty
+[qualifier] test.bigProperty
 with SHORTEN_AND_STAR_IMPORT:
 [qualifier] test.bigProperty
 [qualifier] test.bigProperty
 [qualifier] test.bigProperty
+[qualifier] test.bigProperty

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt

@@ -13,22 +13,22 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLPartialBody
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.partialBodyAnalysisState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirDesignationEntry
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.isAutonomousElement
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FirElementsRecorder.Companion.anchorPsi
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLPartialBodyElementMapper
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.isAutonomousElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.Context
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.ContextKind
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.FilterResponse
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.SessionAndScopeSessionHolder
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
+import org.jetbrains.kotlin.fir.declarations.utils.isScriptTopLevelDeclaration
 import org.jetbrains.kotlin.fir.declarations.utils.memberDeclarationNameOrNull
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.realPsi
-import org.jetbrains.kotlin.fir.SessionAndScopeSessionHolder
-import org.jetbrains.kotlin.fir.declarations.utils.isScriptTopLevelDeclaration
 import org.jetbrains.kotlin.fir.extensions.scriptResolutionHacksComponent
+import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitValue
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
@@ -143,7 +143,7 @@ object ContextCollector {
     private fun partiallyResolveTargetElementIfPossible(
         resolutionFacade: LLResolutionFacade,
         designation: FirDesignation?,
-        targetElement: PsiElement
+        targetElement: PsiElement,
     ): Boolean {
         val declaration = designation?.target?.realPsi as? KtDeclaration ?: return false
 
@@ -291,7 +291,7 @@ private class ContextCollectorVisitor(
             return
         }
 
-        val psi = fir.psi ?: return
+        val psi = fir.anchorPsi ?: return
 
         val key = ContextKey(psi, kind)
         if (key in result) {
@@ -1087,7 +1087,7 @@ private class ContextCollectorVisitor(
     private inner class FilteringVisitor(
         val delegate: FirVisitorVoid,
         val elementsToSkip: Set<FirElement>,
-        val checkIsActive: Boolean
+        val checkIsActive: Boolean,
     ) : FirVisitorVoid() {
         override fun visitElement(element: FirElement) {
             if (checkIsActive && !isActive) {

analysis/low-level-api-fir/testData/contextCollector/enumClass.body.copy.txt

@@ -11,6 +11,35 @@ Tower Data Context:
         Scope: FirPackageMemberScope
     Element 5
         Scope: FirExplicitSimpleImportingScope
+    Element 6
+        Implicit receiver:
+            FirRegularClassSymbol public final companion object Companion : R|kotlin/Any|
+                Type: kotlin.Enum.Companion
+    Element 7
+        Scope: FirNestedClassifierScopeWithSubstitution
+            Classifiers:
+                FirRegularClassSymbol public final companion object Companion : R|kotlin/Any|
+        Static scope owner symbol: FirRegularClassSymbol public abstract class Enum<E : R|kotlin/Enum<E>|> : R|kotlin/Comparable<E>|, R|java/io/Serializable|
+    Element 8
+        Scope: FirNestedClassifierScopeImpl
+            Classifiers:
+                FirRegularClassSymbol public final? inner class InnerClass : R|kotlin/Any|
+                FirRegularClassSymbol public final? class NestedClass : R|kotlin/Any|
+                FirRegularClassSymbol public final? object NestedObject : R|kotlin/Any|
+        Scope: FirNameAwareOnlyCallablesScope
+            Functions
+                FirNamedFunctionSymbol public final static fun valueOf(value: R|kotlin/String|): R|foo/MyEnumClass|
+                FirNamedFunctionSymbol public final static fun values(): R|kotlin/Array<foo/MyEnumClass>|
+            Properties:
+                FirEnumEntrySymbol public final static enum entry ENTRY1: R|foo/MyEnumClass|
+                FirEnumEntrySymbol public final static enum entry ENTRY2: R|foo/MyEnumClass|
+                FirEnumEntrySymbol public final static enum entry ENTRY3: R|foo/MyEnumClass|
+                FirRegularPropertySymbol public final static val entries: R|kotlin/enums/EnumEntries<foo/MyEnumClass>|
+                    public get(): R|kotlin/enums/EnumEntries<foo/MyEnumClass>|
+    Element 9
+        Implicit receiver:
+            FirRegularClassSymbol public final enum class MyEnumClass : R|kotlin/Enum<foo/MyEnumClass>|
+                Type: foo.MyEnumClass
 
 FILE: [ResolvedTo(IMPORTS)] enumClass.kt
     public final [ResolvedTo(STATUS)] enum class MyEnumClass : R|kotlin/Enum<foo/MyEnumClass>| {

analysis/low-level-api-fir/testData/contextCollector/enumClass.body.txt

@@ -11,6 +11,35 @@ Tower Data Context:
         Scope: FirPackageMemberScope
     Element 5
         Scope: FirExplicitSimpleImportingScope
+    Element 6
+        Implicit receiver:
+            FirRegularClassSymbol public final companion object Companion : R|kotlin/Any|
+                Type: kotlin.Enum.Companion
+    Element 7
+        Scope: FirNestedClassifierScopeWithSubstitution
+            Classifiers:
+                FirRegularClassSymbol public final companion object Companion : R|kotlin/Any|
+        Static scope owner symbol: FirRegularClassSymbol public abstract class Enum<E : R|kotlin/Enum<E>|> : R|kotlin/Comparable<E>|, R|java/io/Serializable|
+    Element 8
+        Scope: FirNestedClassifierScopeImpl
+            Classifiers:
+                FirRegularClassSymbol public final? inner class InnerClass : R|kotlin/Any|
+                FirRegularClassSymbol public final? class NestedClass : R|kotlin/Any|
+                FirRegularClassSymbol public final? object NestedObject : R|kotlin/Any|
+        Scope: FirNameAwareOnlyCallablesScope
+            Functions
+                FirNamedFunctionSymbol public final static fun valueOf(value: R|kotlin/String|): R|foo/MyEnumClass|
+                FirNamedFunctionSymbol public final static fun values(): R|kotlin/Array<foo/MyEnumClass>|
+            Properties:
+                FirEnumEntrySymbol public final static enum entry ENTRY1: R|foo/MyEnumClass|
+                FirEnumEntrySymbol public final static enum entry ENTRY2: R|foo/MyEnumClass|
+                FirEnumEntrySymbol public final static enum entry ENTRY3: R|foo/MyEnumClass|
+                FirRegularPropertySymbol public final static val entries: R|kotlin/enums/EnumEntries<foo/MyEnumClass>|
+                    public get(): R|kotlin/enums/EnumEntries<foo/MyEnumClass>|
+    Element 9
+        Implicit receiver:
+            FirRegularClassSymbol public final enum class MyEnumClass : R|kotlin/Enum<foo/MyEnumClass>|
+                Type: foo.MyEnumClass
 
 FILE: [ResolvedTo(IMPORTS)] enumClass.kt
     public final [ResolvedTo(STATUS)] enum class MyEnumClass : R|kotlin/Enum<foo/MyEnumClass>| {

analysis/low-level-api-fir/testData/contextCollector/enumValuePropertyAccessor.body.copy.txt

@@ -40,6 +40,8 @@ Tower Data Context:
         Scope: FirLocalScope
             Properties:
                 FirBackingFieldSymbol private final backing field field: R|kotlin/Int|
+    Element 11
+        Scope: FirLocalScope
 
 FILE: [ResolvedTo(IMPORTS)] enumValuePropertyAccessor.kt
     public final? [ResolvedTo(RAW_FIR)] enum class Direction : R|kotlin/Enum<Direction>| {

analysis/low-level-api-fir/testData/contextCollector/enumValuePropertyAccessor.body.txt

@@ -40,6 +40,8 @@ Tower Data Context:
         Scope: FirLocalScope
             Properties:
                 FirBackingFieldSymbol private final backing field field: R|kotlin/Int|
+    Element 11
+        Scope: FirLocalScope
 
 FILE: [ResolvedTo(IMPORTS)] enumValuePropertyAccessor.kt
     public final [ResolvedTo(STATUS)] enum class Direction : R|kotlin/Enum<Direction>| {