[FIR IDE] Fixed property and parameter creation/replacement and fixed invalid source on annotation type

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt

@@ -15,9 +15,7 @@ import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.builder.Context
-import org.jetbrains.kotlin.fir.builder.extractContractDescriptionIfPossible
-import org.jetbrains.kotlin.fir.builder.generateAccessorsByDelegate
+import org.jetbrains.kotlin.fir.builder.*
 import org.jetbrains.kotlin.fir.contracts.FirContractDescription
 import org.jetbrains.kotlin.fir.contracts.builder.buildRawContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
@@ -1041,10 +1039,12 @@ class DeclarationsConverter(
                 val receiver = delegateExpression?.let {
                     expressionConverter.getAsFirExpression<FirExpression>(it, "Incorrect delegate expression")
                 }
+
                 generateAccessorsByDelegate(
                     delegateBuilder,
-                    classWrapper?.classBuilder,
                     baseModuleData,
+                    classWrapper?.classBuilder?.ownerRegularOrAnonymousObjectSymbol,
+                    classWrapper?.classBuilder?.ownerRegularClassTypeParametersCount,
                     isExtension = false,
                     stubMode = stubMode,
                     receiver = receiver
@@ -1111,8 +1111,9 @@ class DeclarationsConverter(
                     }
                     generateAccessorsByDelegate(
                         delegateBuilder,
-                        classWrapper?.classBuilder,
                         baseModuleData,
+                        classWrapper?.classBuilder?.ownerRegularOrAnonymousObjectSymbol,
+                        classWrapper?.classBuilder?.ownerRegularClassTypeParametersCount,
                         isExtension = receiverType != null,
                         stubMode = stubMode,
                         receiver = receiver

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/LazyBodyRawFirBuilder.kt

@@ -6,50 +6,127 @@
 package org.jetbrains.kotlin.fir.builder
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.containingClass
-import org.jetbrains.kotlin.fir.containingClassAttr
-import org.jetbrains.kotlin.fir.declarations.FirCallableMemberDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.isInner
-import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.psi.*
 
-class RawFirFragmentForLazyBodiesBuilder private constructor(
+data class RawFirReplacement<T : KtElement>(val from: T, val to: T)
+
+class RawFirFragmentForLazyBodiesBuilder<T : KtElement> private constructor(
     session: FirSession,
     baseScopeProvider: FirScopeProvider,
-    private val declaration: KtDeclaration,
+    private val declarationToBuild: KtDeclaration,
+    private val replacement: RawFirReplacement<T>? = null
 ) : RawFirBuilder(session, baseScopeProvider, RawFirBuilderMode.NORMAL) {
 
+    private var replacementApplied = false
+
     companion object {
+        fun elementIsApplicable(element: KtElement) = when (element) {
+            is KtFile, is KtClassInitializer, is KtClassOrObject, is KtObjectLiteralExpression, is KtTypeAlias,
+            is KtNamedFunction, is KtLambdaExpression, is KtAnonymousInitializer, is KtProperty, is KtTypeReference,
+            is KtAnnotationEntry, is KtTypeParameter, is KtTypeProjection, is KtParameter, is KtBlockExpression,
+            is KtSimpleNameExpression, is KtConstantExpression, is KtStringTemplateExpression, is KtReturnExpression,
+            is KtTryExpression, is KtIfExpression, is KtWhenExpression, is KtDoWhileExpression, is KtWhileExpression,
+            is KtForExpression, is KtBreakExpression, is KtContinueExpression, is KtBinaryExpression, is KtBinaryExpressionWithTypeRHS,
+            is KtIsExpression, is KtUnaryExpression, is KtCallExpression, is KtArrayAccessExpression, is KtQualifiedExpression,
+            is KtThisExpression, is KtSuperExpression, is KtParenthesizedExpression, is KtLabeledExpression, is KtAnnotatedExpression,
+            is KtThrowExpression, is KtDestructuringDeclaration, is KtClassLiteralExpression, is KtCallableReferenceExpression,
+            is KtCollectionLiteralExpression -> true
+            else -> false
+        }
+
+        fun <T : KtElement> buildWithReplacement(
+            session: FirSession,
+            baseScopeProvider: FirScopeProvider,
+            designation: List<FirDeclaration>,
+            declarationToBuild: KtDeclaration,
+            replacement: RawFirReplacement<T>? = null
+        ): FirDeclaration {
+            if (replacement != null) {
+                require(elementIsApplicable(replacement.from)) {
+                    "Build with replacement is possible for applicable type but given ${replacement.from::class.simpleName}"
+                }
+                require(replacement.from::class == replacement.to::class) {
+                    "Build with replacement is possible for same type in replacements but given\n${replacement.from::class.simpleName} and ${replacement.to::class.simpleName}"
+                }
+            }
+
+            val builder = RawFirFragmentForLazyBodiesBuilder(session, baseScopeProvider, declarationToBuild, replacement)
+            builder.context.packageFqName = declarationToBuild.containingKtFile.packageFqName
+
+            val result = builder.moveNext(designation.iterator(), containingClass = null)
+            check(replacement == null || builder.replacementApplied) {
+                "Replacement requested but was not applied for ${replacement!!.from::class.simpleName}"
+            }
+            return result
+        }
+
         fun build(
             session: FirSession,
             baseScopeProvider: FirScopeProvider,
             designation: List<FirDeclaration>,
-            declaration: KtDeclaration,
+            rootNonLocalDeclaration: KtDeclaration
         ): FirDeclaration {
-            val builder = RawFirFragmentForLazyBodiesBuilder(session, baseScopeProvider, declaration)
-            builder.context.packageFqName = declaration.containingKtFile.packageFqName
-            return builder.moveNext(designation.iterator())
+            val builder = RawFirFragmentForLazyBodiesBuilder<KtElement>(session, baseScopeProvider, rootNonLocalDeclaration)
+            builder.context.packageFqName = rootNonLocalDeclaration.containingKtFile.packageFqName
+            return builder.moveNext(designation.iterator(), containingClass = null)
         }
     }
 
-    private fun moveNext(iterator: Iterator<FirDeclaration>): FirDeclaration {
+    private fun KtElement.replaced(): KtElement {
+        if (replacement == null || replacement.from != this) return this
+        replacementApplied = true
+        return replacement.to
+    }
+
+    private inner class VisitorWithReplacement : Visitor() {
+        override fun convertElement(element: KtElement): FirElement? =
+            super.convertElement(element.replaced())
+
+        override fun convertProperty(
+            property: KtProperty,
+            ownerRegularOrAnonymousObjectSymbol: FirClassSymbol<*>?,
+            ownerRegularClassTypeParametersCount: Int?
+        ): FirProperty {
+            val replacementProperty = property.replaced()
+            check(replacementProperty is KtProperty)
+            return super.convertProperty(
+                property = replacementProperty,
+                ownerRegularOrAnonymousObjectSymbol = ownerRegularOrAnonymousObjectSymbol,
+                ownerRegularClassTypeParametersCount = ownerRegularClassTypeParametersCount
+            )
+        }
+
+        override fun convertValueParameter(valueParameter: KtParameter, defaultTypeRef: FirTypeRef?): FirValueParameter {
+            val replacementParameter = valueParameter.replaced()
+            check(replacementParameter is KtParameter)
+            return super.convertValueParameter(
+                valueParameter = replacementParameter,
+                defaultTypeRef = defaultTypeRef
+            )
+        }
+    }
+
+    private fun moveNext(iterator: Iterator<FirDeclaration>, containingClass: FirRegularClass?): FirDeclaration {
         if (!iterator.hasNext()) {
-            return if (declaration is KtProperty) {
-                with(Visitor()) {
-                    declaration.toFirProperty(null)
+            val visitor = VisitorWithReplacement()
+            return when (declarationToBuild) {
+                is KtProperty -> {
+                    val ownerSymbol = containingClass?.symbol
+                    val ownerTypeArgumentsCount = containingClass?.typeParameters?.size
+                    visitor.convertProperty(declarationToBuild, ownerSymbol, ownerTypeArgumentsCount)
                 }
-            } else {
-                declaration.accept(Visitor(), Unit) as FirDeclaration
-            }
+                else -> visitor.convertElement(declarationToBuild)
+            } as FirDeclaration
         }
 
         val parent = iterator.next()
-        if (parent !is FirRegularClass) return moveNext(iterator)
+        if (parent !is FirRegularClass) return moveNext(iterator, containingClass = null)
 
         val classOrObject = parent.psi
         check(classOrObject is KtClassOrObject)
@@ -59,7 +136,7 @@ class RawFirFragmentForLazyBodiesBuilder private constructor(
                 if (!parent.isInner) context.capturedTypeParameters = context.capturedTypeParameters.clear()
                 addCapturedTypeParameters(parent.typeParameters.take(classOrObject.typeParameters.size))
                 registerSelfType(classOrObject.toDelegatedSelfType(parent))
-                return moveNext(iterator)
+                return moveNext(iterator, parent)
             }
         }
     }

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -154,12 +154,23 @@ open class RawFirBuilder(
             }
         }
 
-    protected inner class Visitor : KtVisitor<FirElement, Unit>() {
+    protected open inner class Visitor : KtVisitor<FirElement, Unit>() {
         private inline fun <reified R : FirElement> KtElement?.convertSafe(): R? =
-            this?.accept(this@Visitor, Unit) as? R
+            this?.let(::convertElement) as? R
 
         private inline fun <reified R : FirElement> KtElement.convert(): R =
-            this.accept(this@Visitor, Unit) as R
+            convertElement(this) as R
+
+        open fun convertElement(element: KtElement): FirElement? =
+            element.accept(this@Visitor, Unit)
+
+        open fun convertProperty(
+            property: KtProperty, ownerRegularOrAnonymousObjectSymbol: FirClassSymbol<*>?,
+            ownerRegularClassTypeParametersCount: Int?,
+        ): FirProperty = property.toFirProperty(ownerRegularOrAnonymousObjectSymbol, ownerRegularClassTypeParametersCount)
+
+        open fun convertValueParameter(valueParameter: KtParameter, defaultTypeRef: FirTypeRef? = null): FirValueParameter =
+            valueParameter.toFirValueParameter(defaultTypeRef)
 
         private fun KtTypeReference?.toFirOrImplicitType(): FirTypeRef =
             convertSafe() ?: buildImplicitTypeRef {
@@ -230,7 +241,11 @@ open class RawFirBuilder(
                     }
                 }
                 is KtProperty -> {
-                    toFirProperty(ownerClassBuilder)
+                    convertProperty(
+                        this@toFirDeclaration,
+                        ownerClassBuilder.ownerRegularOrAnonymousObjectSymbol,
+                        ownerClassBuilder.ownerRegularClassTypeParametersCount
+                    )
                 }
                 else -> convert()
             }
@@ -521,7 +536,7 @@ open class RawFirBuilder(
             defaultTypeRef: FirTypeRef? = null,
         ) {
             for (valueParameter in valueParameters) {
-                container.valueParameters += valueParameter.toFirValueParameter(defaultTypeRef)
+                container.valueParameters += convertValueParameter(valueParameter, defaultTypeRef)
             }
         }
 
@@ -1120,7 +1135,7 @@ open class RawFirBuilder(
                         val typeRef = valueParameter.typeReference?.convertSafe() ?: buildImplicitTypeRef {
                             source = implicitTypeRefSource
                         }
-                        valueParameter.toFirValueParameter(typeRef)
+                        convertValueParameter(valueParameter, typeRef)
                     }
                 }
                 val expressionSource = expression.toFirSourceElement()
@@ -1237,7 +1252,10 @@ open class RawFirBuilder(
             }
         }
 
-        fun KtProperty.toFirProperty(ownerClassBuilder: FirClassBuilder?): FirProperty {
+        private fun KtProperty.toFirProperty(
+            ownerRegularOrAnonymousObjectSymbol: FirClassSymbol<*>?,
+            ownerRegularClassTypeParametersCount: Int?,
+        ): FirProperty {
             val propertyType = typeReference.toFirOrImplicitType()
             val propertyName = nameAsSafeName
             val isVar = isVar
@@ -1280,8 +1298,9 @@ open class RawFirBuilder(
                     val receiver = delegateExpression?.toFirExpression("Incorrect delegate expression")
                     generateAccessorsByDelegate(
                         delegateBuilder,
-                        null,
                         baseModuleData,
+                        ownerRegularOrAnonymousObjectSymbol = null,
+                        ownerRegularClassTypeParametersCount = null,
                         isExtension = false,
                         stubMode = stubMode,
                         receiver = receiver
@@ -1328,8 +1347,9 @@ open class RawFirBuilder(
                         val receiver = delegateExpression?.toFirExpression("Should have delegate")
                         generateAccessorsByDelegate(
                             delegateBuilder,
-                            ownerClassBuilder,
                             baseModuleData,
+                            ownerRegularOrAnonymousObjectSymbol,
+                            ownerRegularClassTypeParametersCount,
                             isExtension = receiverTypeReference != null,
                             stubMode = stubMode,
                             receiver = receiver
@@ -1354,7 +1374,7 @@ open class RawFirBuilder(
         }
 
         override fun visitProperty(property: KtProperty, data: Unit): FirElement {
-            return property.toFirProperty(ownerClassBuilder = null)
+            return property.toFirProperty(ownerRegularOrAnonymousObjectSymbol = null, ownerRegularClassTypeParametersCount = null)
         }
 
         override fun visitTypeReference(typeReference: KtTypeReference, data: Unit): FirElement {
@@ -1534,7 +1554,7 @@ open class RawFirBuilder(
         }
 
         override fun visitParameter(parameter: KtParameter, data: Unit): FirElement =
-            parameter.toFirValueParameter()
+            convertValueParameter(parameter)
 
         override fun visitBlockExpression(expression: KtBlockExpression, data: Unit): FirElement {
             return configureBlockWithoutBuilding(expression).build()
@@ -1597,7 +1617,7 @@ open class RawFirBuilder(
                 tryBlock = expression.tryBlock.toFirBlock()
                 finallyBlock = expression.finallyBlock?.finalExpression?.toFirBlock()
                 for (clause in expression.catchClauses) {
-                    val parameter = clause.catchParameter?.toFirValueParameter() ?: continue
+                    val parameter = clause.catchParameter?.let { convertValueParameter(it) } ?: continue
                     catches += buildCatch {
                         source = clause.toFirSourceElement()
                         this.parameter = parameter

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -28,10 +28,7 @@ import org.jetbrains.kotlin.fir.references.builder.buildImplicitThisReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildSimpleNamedReference
 import org.jetbrains.kotlin.fir.symbols.constructStarProjectedType
-import org.jetbrains.kotlin.fir.symbols.impl.FirDelegateFieldSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirPropertyAccessorSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.*
 import org.jetbrains.kotlin.fir.types.impl.*
@@ -293,10 +290,21 @@ fun generateTemporaryVariable(
     moduleData: FirModuleData, source: FirSourceElement?, specialName: String, initializer: FirExpression,
 ): FirVariable<*> = generateTemporaryVariable(moduleData, source, Name.special("<$specialName>"), initializer)
 
+val FirClassBuilder.ownerRegularOrAnonymousObjectSymbol
+    get() = when (this) {
+        is FirAnonymousObjectBuilder -> symbol
+        is FirRegularClassBuilder -> symbol
+        else -> null
+    }
+
+val FirClassBuilder.ownerRegularClassTypeParametersCount
+    get() = if (this is FirRegularClassBuilder) typeParameters.size else null
+
 fun FirPropertyBuilder.generateAccessorsByDelegate(
     delegateBuilder: FirWrappedDelegateExpressionBuilder?,
-    ownerClassBuilder: FirClassBuilder?,
     moduleData: FirModuleData,
+    ownerRegularOrAnonymousObjectSymbol: FirClassSymbol<*>?,
+    ownerRegularClassTypeParametersCount: Int?,
     isExtension: Boolean,
     stubMode: Boolean,
     receiver: FirExpression?
@@ -305,12 +313,8 @@ fun FirPropertyBuilder.generateAccessorsByDelegate(
     val delegateFieldSymbol = FirDelegateFieldSymbol<FirProperty>(symbol.callableId).also {
         this.delegateFieldSymbol = it
     }
-    val ownerSymbol = when (ownerClassBuilder) {
-        is FirAnonymousObjectBuilder -> ownerClassBuilder.symbol
-        is FirRegularClassBuilder -> ownerClassBuilder.symbol
-        else -> null
-    }
-    val isMember = ownerSymbol != null
+
+    val isMember = ownerRegularOrAnonymousObjectSymbol != null
     val fakeSource = delegateBuilder.source?.fakeElement(FirFakeSourceElementKind.DelegatedPropertyAccessor)
 
     /*
@@ -335,14 +339,14 @@ fun FirPropertyBuilder.generateAccessorsByDelegate(
                     boundSymbol = this@generateAccessorsByDelegate.symbol
                 }
             }
-            ownerSymbol != null -> buildThisReceiverExpression {
+            ownerRegularOrAnonymousObjectSymbol != null -> buildThisReceiverExpression {
                 source = fakeSource
                 calleeReference = buildImplicitThisReference {
-                    boundSymbol = ownerSymbol
+                    boundSymbol = ownerRegularOrAnonymousObjectSymbol
                 }
                 typeRef = buildResolvedTypeRef {
-                    val typeParameterNumber = (ownerClassBuilder as? FirRegularClassBuilder)?.typeParameters?.size ?: 0
-                    type = ownerSymbol.constructStarProjectedType(typeParameterNumber)
+                    val typeParameterNumber = ownerRegularClassTypeParametersCount ?: 0
+                    type = ownerRegularOrAnonymousObjectSymbol.constructStarProjectedType(typeParameterNumber)
                 }
             }
             else -> buildConstExpression(null, ConstantValueKind.Null, null)
@@ -353,7 +357,7 @@ fun FirPropertyBuilder.generateAccessorsByDelegate(
         calleeReference = buildDelegateFieldReference {
             resolvedSymbol = delegateFieldSymbol
         }
-        if (ownerSymbol != null) {
+        if (ownerRegularOrAnonymousObjectSymbol != null) {
             dispatchReceiver = thisRef(forDispatchReceiver = true)
         }
     }