FIR IDE: fix ref shortener on property in primary constructor

Currently the shortener ignores the created FirProperty, as a result,
any annotations declared on the getter or setter are ignored. The test
is added in Intellij repo.

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

@@ -17,26 +17,18 @@ import org.jetbrains.kotlin.analysis.api.fir.utils.computeImportableName
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.*
-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.getOrBuildFirOfType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.descriptors.ClassKind
-import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE
 import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
 import org.jetbrains.kotlin.fir.declarations.builder.buildImport
 import org.jetbrains.kotlin.fir.declarations.builder.buildResolvedImport
 import org.jetbrains.kotlin.fir.expressions.FirErrorResolvedQualifier
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
-import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
@@ -90,7 +82,9 @@ internal class KtFirReferenceShortener(
             towerContext,
             selection,
             classShortenOption = { classShortenOption(analysisSession.firSymbolBuilder.buildSymbol(it.fir) as KtClassLikeSymbol) },
-            callableShortenOption = { callableShortenOption(analysisSession.firSymbolBuilder.buildSymbol(it.fir) as KtCallableSymbol) })
+            callableShortenOption = { callableShortenOption(analysisSession.firSymbolBuilder.buildSymbol(it.fir) as KtCallableSymbol) },
+            firResolveState,
+        )
         firDeclaration.accept(collector)
 
         return ShortenCommandImpl(
@@ -279,12 +273,41 @@ private class ElementsToShortenCollector(
     private val selection: TextRange,
     private val classShortenOption: (FirClassLikeSymbol<*>) -> ShortenOption,
     private val callableShortenOption: (FirCallableSymbol<*>) -> ShortenOption,
+    private val firResolveState: FirModuleResolveState,
 ) :
     FirVisitorVoid() {
     val namesToImport: MutableList<FqName> = mutableListOf()
     val namesToImportWithStar: MutableList<FqName> = mutableListOf()
     val typesToShorten: MutableList<KtUserType> = mutableListOf()
     val qualifiersToShorten: MutableList<KtDotQualifiedExpression> = mutableListOf()
+    private val visitedProperty = mutableSetOf<FirProperty>()
+
+    override fun visitValueParameter(valueParameter: FirValueParameter) {
+        super.visitValueParameter(valueParameter)
+        valueParameter.getCorrespondingProperty()?.let { visitProperty(it) }
+    }
+
+    private fun FirValueParameter.getCorrespondingProperty(): FirProperty? {
+        val parameter = psi as? KtParameter
+        val containingFunction = parameter?.getContainingFunction()
+        if (containingFunction is KtPrimaryConstructor && parameter.valOrVarKeyword != null) {
+            // also check the FirProperty corresponding to this parameter in the primary constructor
+            return (containingFunction.parent as? KtClass)
+                ?.getOrBuildFirSafe<FirClass>(firResolveState)
+                ?.declarations
+                ?.firstOrNull { it.source?.kind == FirFakeSourceElementKind.PropertyFromParameter && it.psi == parameter } as FirProperty
+        }
+        return null
+    }
+
+    private fun KtParameter.getContainingFunction(): KtFunction? = (parent as? KtParameterList)?.parent as? KtFunction
+
+
+    override fun visitProperty(property: FirProperty) {
+        if (visitedProperty.add(property)) {
+            super.visitProperty(property)
+        }
+    }
 
     override fun visitElement(element: FirElement) {
         element.acceptChildren(this)