IDE FIR: fix resolving vararg parameter type

For a vararg parameter type, there corresponding FIR element has a fake
source of kind ArrayTypeFromVarargParameter. As a result,
`getOrBuildFir` returns the whole `FirValueParameter` for the parameter
type reference. Therefore, we need some special handling for this case
in order to resolve the proper `KtSymbol`.

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/TransformUtils.kt

@@ -9,10 +9,15 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.declarations.FirTypedDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
-import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.visitors.*
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
 
 internal object StoreType : FirDefaultTransformer<FirTypeRef>() {
     override fun <E : FirElement> transformElement(element: E, data: FirTypeRef): CompositeTransformResult<E> {
@@ -96,13 +101,30 @@ internal fun FirValueParameter.transformVarargTypeToArrayType() {
 }
 
 internal fun FirTypedDeclaration.transformTypeToArrayType() {
+    val returnTypeRef = this.returnTypeRef
+    require(returnTypeRef is FirResolvedTypeRef)
+    // If the delegated type is already resolved, it means we have already created a resolved array type for this vararg type declaration.
+    // This is because in the buildResolvedTypeRef call below, we set the delegated type ref to the previous (non-vararg) resolved type ref.
+    if (returnTypeRef.delegatedTypeRef is FirResolvedTypeRef &&
+        returnTypeRef.delegatedTypeRef?.source?.kind == FirFakeSourceElementKind.ArrayTypeFromVarargParameter
+    ) return
     val returnType = returnTypeRef.coneType
+
     transformReturnTypeRef(
         StoreType,
-        returnTypeRef.withReplacedConeType(
-            ConeKotlinTypeProjectionOut(returnType).createArrayType(),
-            FirFakeSourceElementKind.ArrayTypeFromVarargParameter
-        )
+        buildResolvedTypeRef {
+            source = returnTypeRef.source
+            type = ConeKotlinTypeProjectionOut(returnType).createArrayType()
+            annotations += returnTypeRef.annotations
+            delegatedTypeRef = returnTypeRef.apply {
+                // Here we replace the source of the immediate delegate and nested delegate with a fake source. Normally the immediate
+                // delegate is the original resolved type ref before this array augmentation. That is, this is the array element type.
+                // The nested delegate is the raw type ref created when constructing the FIR. Usually it's a `FirUserTypeRef` for value
+                // parameters.
+                replaceSource(source?.fakeElement(FirFakeSourceElementKind.ArrayTypeFromVarargParameter))
+                delegatedTypeRef?.replaceSource(source?.fakeElement(FirFakeSourceElementKind.ArrayTypeFromVarargParameter))
+            }
+        }
     )
 }
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirUserTypeRefImpl.kt

@@ -16,7 +16,7 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.fir.visitors.transformInplace
 
 class FirUserTypeRefImpl(
-    override val source: FirSourceElement?,
+    override var source: FirSourceElement?,
     override val isMarkedNullable: Boolean,
     override val qualifier: MutableList<FirQualifierPart>,
     override val annotations: MutableList<FirAnnotationCall>
@@ -43,5 +43,6 @@ class FirUserTypeRefImpl(
     }
 
     override fun replaceSource(newSource: FirSourceElement?) {
+        source = newSource
     }
-}
+}