[FIR] During resolve, set correctly property reference type in delegate

compiler/fir/analysis-tests/testData/resolveWithStdlib/delegates/simpleDelegatedToMap.kt

@@ -1,5 +1,6 @@
 class C(val map: MutableMap<String, Any>) {
-    var foo by map
+    // NB: this does not work because of @LowPriorityInOverloadResolution not deserialized (KT-37228)
+    var foo by <!AMBIGUITY!>map<!>
 }
 
-var bar by hashMapOf<String, Any>()
+var bar by <!AMBIGUITY!>hashMapOf<String, Any>()<!>

compiler/fir/cones/src/org/jetbrains/kotlin/fir/symbols/StandardClassIds.kt

@@ -42,6 +42,12 @@ object StandardClassIds {
     val String = "String".baseId()
 
     val KProperty = "KProperty".reflectId()
+    val KProperty0 = "KProperty0".reflectId()
+    val KMutableProperty0 = "KMutableProperty0".reflectId()
+    val KProperty1 = "KProperty1".reflectId()
+    val KMutableProperty1 = "KMutableProperty1".reflectId()
+    val KProperty2 = "KProperty2".reflectId()
+    val KMutableProperty2 = "KMutableProperty2".reflectId()
 
     val Comparable = "Comparable".baseId()
     val Number = "Number".baseId()

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

@@ -36,8 +36,7 @@ import org.jetbrains.kotlin.fir.types.FirUserTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildTypeProjectionWithVariance
 import org.jetbrains.kotlin.fir.types.builder.buildUserTypeRef
-import org.jetbrains.kotlin.fir.types.impl.FirImplicitKPropertyTypeRef
-import org.jetbrains.kotlin.fir.types.impl.FirQualifierPartImpl
+import org.jetbrains.kotlin.fir.types.impl.*
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
@@ -325,6 +324,7 @@ fun FirPropertyBuilder.generateAccessorsByDelegate(
         }
         else buildConstExpression(null, FirConstKind.Null, null)
 
+    val isVar = this@generateAccessorsByDelegate.isVar
     fun propertyRef() = buildCallableReferenceAccess {
         source = delegateBuilder.source
         calleeReference = buildResolvedNamedReference {
@@ -332,7 +332,23 @@ fun FirPropertyBuilder.generateAccessorsByDelegate(
             name = this@generateAccessorsByDelegate.name
             resolvedSymbol = this@generateAccessorsByDelegate.symbol
         }
-        typeRef = FirImplicitKPropertyTypeRef(null, ConeStarProjection)
+        typeRef = when {
+            !member && !extension -> if (isVar) {
+                FirImplicitKMutableProperty0TypeRef(null, ConeStarProjection)
+            } else {
+                FirImplicitKProperty0TypeRef(null, ConeStarProjection)
+            }
+            member && extension -> if (isVar) {
+                FirImplicitKMutableProperty2TypeRef(null, ConeStarProjection, ConeStarProjection, ConeStarProjection)
+            } else {
+                FirImplicitKProperty2TypeRef(null, ConeStarProjection, ConeStarProjection, ConeStarProjection)
+            }
+            else -> if (isVar) {
+                FirImplicitKMutableProperty1TypeRef(null, ConeStarProjection, ConeStarProjection)
+            } else {
+                FirImplicitKProperty1TypeRef(null, ConeStarProjection, ConeStarProjection)
+            }
+        }
     }
 
     delegateBuilder.delegateProvider = if (stubMode) buildExpressionStub() else buildFunctionCall {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.visitors.*
 import org.jetbrains.kotlin.name.Name
 
@@ -139,6 +140,45 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
         }
     }
 
+    private fun FirFunctionCall.replacePropertyReferenceTypeInDelegateAccessors(property: FirProperty) {
+        // var someProperty: SomeType
+        //     get() = delegate.getValue(thisRef, kProperty: KProperty0/1/2<..., SomeType>)
+        //     set() = delegate.getValue(thisRef, kProperty: KProperty0/1/2<..., SomeType>, value)
+        val propertyReferenceAccess = argumentMapping?.keys?.toList()?.getOrNull(1) as? FirCallableReferenceAccess ?: return
+        val typeRef = propertyReferenceAccess.typeRef
+        if (typeRef is FirResolvedTypeRef && property.returnTypeRef is FirResolvedTypeRef) {
+            val typeArguments = (typeRef.type as ConeClassLikeType).typeArguments
+            val extensionType = property.receiverTypeRef?.coneTypeSafe<ConeKotlinType>()
+            propertyReferenceAccess.replaceTypeRef(
+                buildResolvedTypeRef {
+                    source = typeRef.source
+                    annotations.addAll(typeRef.annotations)
+                    type = (typeRef.type as ConeClassLikeType).lookupTag.constructClassType(
+                        typeArguments.mapIndexed { index, argument ->
+                            when (index) {
+                                typeArguments.lastIndex -> property.returnTypeRef.coneTypeUnsafe()
+                                0 -> containingClass?.let { containingClass ->
+                                    containingClass.symbol.constructType(
+                                        Array(containingClass.typeParameters.size) { ConeStarProjection }, isNullable = false
+                                    )
+                                } ?: extensionType
+                                else -> extensionType
+                            } ?: argument
+                        }.toTypedArray(),
+                        isNullable = false
+                    )
+                }
+            )
+        }
+    }
+
+    private fun replacePropertyReferenceTypeInDelegateAccessors(property: FirProperty) {
+        (property.getter?.body?.statements?.singleOrNull() as? FirReturnExpression)?.let { returnExpression ->
+            (returnExpression.result as? FirFunctionCall)?.replacePropertyReferenceTypeInDelegateAccessors(property)
+        }
+        (property.setter?.body?.statements?.singleOrNull() as? FirFunctionCall)?.replacePropertyReferenceTypeInDelegateAccessors(property)
+    }
+
     private fun transformPropertyWithDelegate(property: FirProperty) {
         property.transformDelegate(transformer, ResolutionMode.ContextDependent)
 
@@ -165,6 +205,9 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
             val declarationCompletionResultsWriter = FirDeclarationCompletionResultsWriter(finalSubstitutor)
             property.transformSingle(declarationCompletionResultsWriter, null)
         }
+        if (property.delegateFieldSymbol != null) {
+            replacePropertyReferenceTypeInDelegateAccessors(property)
+        }
         property.transformOtherChildren(transformer, ResolutionMode.ContextIndependent)
     }
 

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

@@ -87,3 +87,45 @@ class FirImplicitKPropertyTypeRef(
     source: FirSourceElement?,
     typeArgument: ConeTypeProjection
 ) : FirImplicitBuiltinTypeRef(source, StandardClassIds.KProperty, arrayOf(typeArgument))
+
+class FirImplicitKProperty0TypeRef(
+    source: FirSourceElement?,
+    propertyTypeArgument: ConeTypeProjection
+) : FirImplicitBuiltinTypeRef(source, StandardClassIds.KProperty0, arrayOf(propertyTypeArgument))
+
+class FirImplicitKMutableProperty0TypeRef(
+    source: FirSourceElement?,
+    propertyTypeArgument: ConeTypeProjection
+) : FirImplicitBuiltinTypeRef(source, StandardClassIds.KMutableProperty0, arrayOf(propertyTypeArgument))
+
+class FirImplicitKProperty1TypeRef(
+    source: FirSourceElement?,
+    receiverTypeArgument: ConeTypeProjection,
+    propertyTypeArgument: ConeTypeProjection
+) : FirImplicitBuiltinTypeRef(source, StandardClassIds.KProperty1, arrayOf(receiverTypeArgument, propertyTypeArgument))
+
+class FirImplicitKMutableProperty1TypeRef(
+    source: FirSourceElement?,
+    receiverTypeArgument: ConeTypeProjection,
+    propertyTypeArgument: ConeTypeProjection
+) : FirImplicitBuiltinTypeRef(source, StandardClassIds.KMutableProperty1, arrayOf(receiverTypeArgument, propertyTypeArgument))
+
+class FirImplicitKProperty2TypeRef(
+    source: FirSourceElement?,
+    dispatchReceiverTypeArgument: ConeTypeProjection,
+    extensionReceiverTypeArgument: ConeTypeProjection,
+    propertyTypeArgument: ConeTypeProjection
+) : FirImplicitBuiltinTypeRef(
+    source, StandardClassIds.KProperty2,
+    arrayOf(dispatchReceiverTypeArgument, extensionReceiverTypeArgument, propertyTypeArgument)
+)
+
+class FirImplicitKMutableProperty2TypeRef(
+    source: FirSourceElement?,
+    dispatchReceiverTypeArgument: ConeTypeProjection,
+    extensionReceiverTypeArgument: ConeTypeProjection,
+    propertyTypeArgument: ConeTypeProjection
+) : FirImplicitBuiltinTypeRef(
+    source, StandardClassIds.KMutableProperty2,
+    arrayOf(dispatchReceiverTypeArgument, extensionReceiverTypeArgument, propertyTypeArgument)
+)