KT-61889 [AA] Migrate KtFirReferenceShortener to ContextCollector

This should make reference shortener considerably faster, since it won't
need to perform redundant extra resolve of the file.

`ContextCollector` more accurately collects the scopes for the scripts,
so some script tests are also fixed.

It should fix the following bugs:

^KTIJ-26714 Fixed
^KTIJ-26727 Fixed

This is also an important part of fixing the following bugs:
- KTIJ-26715
- KTIJ-26734

But those bugs also rely on KT-61890, because completion uses scopes
and snows incorrect elements from them

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt

@@ -26,9 +26,11 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 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.getOrBuildFir
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.AllCandidatesResolver
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector
 import org.jetbrains.kotlin.analysis.utils.printer.parentsOfType
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
@@ -43,6 +45,7 @@ import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildSimpleNamedReference
+import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeAmbiguityError
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnmatchedTypeArgumentsError
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
@@ -94,11 +97,7 @@ internal class KtFirReferenceShortener(
             kDocQualifiersToShorten = emptyList(),
         )
 
-        val towerContext = if (declarationToVisit is KtFile) {
-            LowLevelFirApiFacadeForResolveOnAir.getOnAirTowerDataContextProviderForTheWholeFile(firResolveSession, declarationToVisit)
-        } else {
-            LowLevelFirApiFacadeForResolveOnAir.getOnAirGetTowerContextProvider(firResolveSession, declarationToVisit)
-        }
+        val towerContext = createTowerDataContextProvider(declarationToVisit)
 
         //TODO: collect all usages of available symbols in the file and prevent importing symbols that could introduce name clashes, which
         // may alter the meaning of existing code.
@@ -140,6 +139,10 @@ internal class KtFirReferenceShortener(
         )
     }
 
+    private fun createTowerDataContextProvider(declaration: KtElement): FirTowerContextProvider {
+        return FirTowerContextProviderByContextCollector(declaration, firResolveSession)
+    }
+
     private fun KtElement.getCorrespondingFirElement(): FirElement? {
         require(this is KtFile || this is KtDeclaration)
 
@@ -156,6 +159,45 @@ internal class KtFirReferenceShortener(
     private fun buildSymbol(firSymbol: FirBasedSymbol<*>): KtSymbol = analysisSession.firSymbolBuilder.buildSymbol(firSymbol)
 }
 
+private class FirTowerContextProviderByContextCollector(
+    targetElement: KtElement,
+    firResolveSession: LLFirResolveSession,
+) : FirTowerContextProvider {
+
+    private val contextProvider: ContextCollector.ContextProvider = run {
+        val firFile = targetElement.containingKtFile.getOrBuildFirFile(firResolveSession)
+
+        val sessionHolder = run {
+            val firSession = firResolveSession.useSiteFirSession
+            val scopeSession = firResolveSession.getScopeSessionFor(firSession)
+
+            SessionHolderImpl(firSession, scopeSession)
+        }
+
+        val designation = ContextCollector.computeDesignation(firFile, targetElement)
+
+        ContextCollector.process(
+            firFile,
+            sessionHolder,
+            designation,
+            shouldCollectBodyContext = false, // we only query SELF context
+            filter = { ContextCollector.FilterResponse.CONTINUE }
+        )
+    }
+
+    override fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext? {
+        for (parent in ktElement.parentsWithSelf) {
+            val context = contextProvider[parent, ContextCollector.ContextKind.SELF]
+
+            if (context != null) {
+                return context.towerDataContext
+            }
+        }
+
+        return null
+    }
+}
+
 private fun FqName.dropFakeRootPrefixIfPresent(): FqName =
     tail(FqName(ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE))
 

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/references/FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated.java

@@ -640,6 +640,18 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
                 runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnParameter_nested.kt");
             }
 
+            @Test
+            @TestMetadata("constructorParameterVsTopLevelProperty_conflict.kt")
+            public void testConstructorParameterVsTopLevelProperty_conflict() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/constructorParameterVsTopLevelProperty_conflict.kt");
+            }
+
+            @Test
+            @TestMetadata("constructorParameterVsTopLevelProperty_noConflict.kt")
+            public void testConstructorParameterVsTopLevelProperty_noConflict() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/constructorParameterVsTopLevelProperty_noConflict.kt");
+            }
+
             @Test
             @TestMetadata("contextReceiver.kt")
             public void testContextReceiver() throws Exception {

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/references/FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenerated.java

@@ -640,6 +640,18 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenera
                 runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnParameter_nested.kt");
             }
 
+            @Test
+            @TestMetadata("constructorParameterVsTopLevelProperty_conflict.kt")
+            public void testConstructorParameterVsTopLevelProperty_conflict() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/constructorParameterVsTopLevelProperty_conflict.kt");
+            }
+
+            @Test
+            @TestMetadata("constructorParameterVsTopLevelProperty_noConflict.kt")
+            public void testConstructorParameterVsTopLevelProperty_noConflict() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/constructorParameterVsTopLevelProperty_noConflict.kt");
+            }
+
             @Test
             @TestMetadata("contextReceiver.kt")
             public void testContextReceiver() throws Exception {

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/callInsideScriptExpression.txt

@@ -1,6 +1,7 @@
 Before shortening: one.A
 with DO_NOT_SHORTEN:
 with SHORTEN_IF_ALREADY_IMPORTED:
+[qualifier] one.A
 with SHORTEN_AND_IMPORT:
 [qualifier] one.A
 with SHORTEN_AND_STAR_IMPORT:

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/constructorParameterVsTopLevelProperty_conflict.kt

@@ -0,0 +1,24 @@
+package test
+
+operator fun String.getValue(a: Any, b: Any): String = this
+
+val bigProperty: String = "Hello"
+
+open class Base(s: String) {}
+
+<expr>
+class Foo(
+    bigProperty: String,
+    otherParam: String = test.bigProperty
+) : Base(test.bigProperty),
+    CharSequence by test.bigProperty
+{
+    val regularProperty = test.bigProperty
+    val delegatedProperty by test.bigProperty
+    val lambdaCapture = { test.bigProperty }
+
+    init {
+        test.bigProperty()
+    }
+}
+</expr>

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/constructorParameterVsTopLevelProperty_conflict.txt

@@ -0,0 +1,18 @@
+Before shortening: class Foo(
+    bigProperty: String,
+    otherParam: String = test.bigProperty
+) : Base(test.bigProperty),
+    CharSequence by test.bigProperty
+{
+    val regularProperty = test.bigProperty
+    val delegatedProperty by test.bigProperty
+    val lambdaCapture = { test.bigProperty }
+
+    init {
+        test.bigProperty()
+    }
+}
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+with SHORTEN_AND_STAR_IMPORT:

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/constructorParameterVsTopLevelProperty_noConflict.kt

@@ -0,0 +1,18 @@
+package test
+
+val bigProperty: String = "Hello"
+
+<expr>
+class Foo(bigProperty: String) {
+    val propertyWithAccessors
+        get() = test.bigProperty
+        set(value) {
+            test.bigProperty
+        }
+
+    fun functionWithInitializer() = test.bigProperty
+    fun functionWithBody() {
+        return test.bigProperty
+    }
+}
+</expr>

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

@@ -0,0 +1,28 @@
+Before shortening: class Foo(bigProperty: String) {
+    val propertyWithAccessors
+        get() = test.bigProperty
+        set(value) {
+            test.bigProperty
+        }
+
+    fun functionWithInitializer() = test.bigProperty
+    fun functionWithBody() {
+        return 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/analysis-api/testData/components/referenceShortener/referenceShortenerWholeFile/callInsideScriptExpression.txt

@@ -1,5 +1,7 @@
 with DO_NOT_SHORTEN:
 with SHORTEN_IF_ALREADY_IMPORTED:
+[qualifier] one.A
+[qualifier] one.B
 with SHORTEN_AND_IMPORT:
 [qualifier] one.A
 [qualifier] one.B

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

@@ -39,7 +39,7 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.utils.yieldIfNotNull
 
-internal object ContextCollector {
+object ContextCollector {
     enum class ContextKind {
         /** Represents the context of the declaration itself. */
         SELF,
@@ -79,7 +79,7 @@ internal object ContextCollector {
         val isBodyContextCollected = bodyElement != null
         val acceptedElements = targetElement.parentsWithSelf.toSet()
 
-        val contextProvider = process(file, computeDesignation(file, targetElement), holder, isBodyContextCollected) { candidate ->
+        val contextProvider = process(file, holder, computeDesignation(file, targetElement), isBodyContextCollected) { candidate ->
             when (candidate) {
                 targetElement -> FilterResponse.STOP
                 in acceptedElements -> FilterResponse.CONTINUE
@@ -104,7 +104,7 @@ internal object ContextCollector {
         return null
     }
 
-    private fun computeDesignation(file: FirFile, targetElement: PsiElement): FirDesignation? {
+    fun computeDesignation(file: FirFile, targetElement: PsiElement): FirDesignation? {
         val contextKtDeclaration = targetElement.getNonLocalContainingOrThisDeclaration()
         if (contextKtDeclaration != null) {
             val designationPath = FirElementFinder.collectDesignationPath(file, contextKtDeclaration)
@@ -126,15 +126,15 @@ internal object ContextCollector {
      * Processes the [FirFile], collecting contexts for elements matching the [filter].
      *
      * @param file The file to process.
-     * @param designation The declaration to process. If `null`, all declarations in the [file] are processed.
      * @param holder The [SessionHolder] for the session that owns a [file].
+     * @param designation The declaration to process. If `null`, all declarations in the [file] are processed.
      * @param shouldCollectBodyContext If `true`, [ContextKind.BODY] is collected where available.
      * @param filter The filter predicate. Context is collected only for [PsiElement]s for which the [filter] returns `true`.
      */
     fun process(
         file: FirFile,
-        designation: FirDesignation?,
         holder: SessionHolder,
+        designation: FirDesignation?,
         shouldCollectBodyContext: Boolean,
         filter: (PsiElement) -> FilterResponse
     ): ContextProvider {