From 649bf92d98f35aaa9deaffc16aa4cf6386598598 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Wed, 2 Nov 2022 15:46:08 +0100
Subject: [PATCH] [FIR] simplify receiver parameter creation

^KT-54417
---
 .../converter/DeclarationsConverter.kt        | 13 ++++++---
 .../kotlin/fir/builder/RawFirBuilder.kt       | 27 ++++++++++++++++---
 .../kotlin/fir/builder/ConversionUtils.kt     |  6 ++---
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt
index 4d0365413f4..10fa9f3e50b 100644
--- a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt
+++ b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt
@@ -1183,7 +1183,7 @@ class DeclarationsConverter(
                 )
             } else {
                 this.isLocal = false
-                receiverParameter = receiverType?.convertToReceiverParameter()
+                receiverParameter = receiverType?.asReceiverParameter(receiverType?.extractReceiverAnnotations())
 
                 dispatchReceiverType = currentDispatchReceiverType()
                 withCapturedTypeParameters(true, propertySource, firTypeParameters) {
@@ -1611,7 +1611,7 @@ class DeclarationsConverter(
             functionSymbol = FirAnonymousFunctionSymbol()
             FirAnonymousFunctionBuilder().apply {
                 source = functionSource
-                receiverParameter = receiverType?.convertToReceiverParameter()
+                receiverParameter = receiverType?.asReceiverParameter(receiverType?.extractReceiverAnnotations())
                 symbol = functionSymbol
                 isLambda = false
                 hasExplicitParameterList = true
@@ -1626,7 +1626,7 @@ class DeclarationsConverter(
             functionSymbol = FirNamedFunctionSymbol(callableIdForName(functionName))
             FirSimpleFunctionBuilder().apply {
                 source = functionSource
-                receiverParameter = receiverType?.convertToReceiverParameter()
+                receiverParameter = receiverType?.asReceiverParameter(receiverType?.extractReceiverAnnotations())
                 name = functionName
                 status = FirDeclarationStatusImpl(
                     if (isLocal) Visibilities.Local else modifiers.getVisibility(),
@@ -1950,6 +1950,13 @@ class DeclarationsConverter(
         }
     }
 
+    private fun FirTypeRef.extractReceiverAnnotations(): List<FirAnnotationCall> {
+        @Suppress("UNCHECKED_CAST")
+        val receiverAnnotations = annotations.filter { it is FirAnnotationCall && it.useSiteTarget == RECEIVER } as List<FirAnnotationCall>
+        (annotations as? MutableList<FirAnnotation>)?.removeIf { it.useSiteTarget == RECEIVER }
+        return receiverAnnotations
+    }
+
     /**
      * @see org.jetbrains.kotlin.parsing.KotlinParsing.parseTypeRef
      */
diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
index 4e8149c7767..9577a3051b8 100644
--- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
+++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
@@ -189,6 +189,12 @@ open class RawFirBuilder(
                 source = this@toFirOrImplicitType?.toFirSourceElement(KtFakeSourceElementKind.ImplicitTypeRef)
             }
 
+        private fun KtTypeReference.receiverAnnotations(): List<FirAnnotationCall> {
+            return annotationEntries
+                .filter { it.useSiteTarget?.getAnnotationUseSiteTarget() == RECEIVER }
+                .map { it.convert<FirAnnotationCall>() }
+        }
+
         private fun KtTypeReference?.toFirOrUnitType(): FirTypeRef =
             convertSafe() ?: implicitUnitType
 
@@ -1295,13 +1301,19 @@ open class RawFirBuilder(
             } else {
                 typeReference.toFirOrImplicitType()
             }
+
             val receiverType = function.receiverTypeReference.convertSafe<FirTypeRef>()
+            val receiverAnnotations = if (receiverType != null) {
+                function.receiverTypeReference?.receiverAnnotations()
+            } else {
+                null
+            }
 
             val labelName: String?
             val isAnonymousFunction = function.name == null && !function.parent.let { it is KtFile || it is KtClassBody }
             val functionBuilder = if (isAnonymousFunction) {
                 FirAnonymousFunctionBuilder().apply {
-                    receiverParameter = receiverType?.convertToReceiverParameter()
+                    receiverParameter = receiverType?.asReceiverParameter(receiverAnnotations)
                     symbol = FirAnonymousFunctionSymbol()
                     isLambda = false
                     hasExplicitParameterList = true
@@ -1310,7 +1322,7 @@ open class RawFirBuilder(
                 }
             } else {
                 FirSimpleFunctionBuilder().apply {
-                    receiverParameter = receiverType?.convertToReceiverParameter()
+                    receiverParameter = receiverType?.asReceiverParameter(receiverAnnotations)
                     name = function.nameAsSafeName
                     labelName = runIf(!name.isSpecial) { name.identifier }
                     symbol = FirNamedFunctionSymbol(callableIdForName(function.nameAsSafeName))
@@ -1642,7 +1654,14 @@ open class RawFirBuilder(
                     }
                 } else {
                     isLocal = false
-                    receiverParameter = receiverTypeReference.convertSafe<FirTypeRef>()?.convertToReceiverParameter()
+                    val receiverType = receiverTypeReference.convertSafe<FirTypeRef>()
+                    val receiverAnnotations = if (receiverType != null) {
+                        receiverTypeReference?.receiverAnnotations()
+                    } else {
+                        null
+                    }
+
+                    receiverParameter = receiverType?.asReceiverParameter(receiverAnnotations)
                     symbol = FirPropertySymbol(callableIdForName(propertyName))
                     dispatchReceiverType = currentDispatchReceiverType()
                     extractTypeParametersTo(this, symbol)
@@ -1840,9 +1859,11 @@ open class RawFirBuilder(
 
             for (modifierList in allModifierLists) {
                 for (annotationEntry in modifierList.annotationEntries) {
+                    if (annotationEntry.useSiteTarget?.getAnnotationUseSiteTarget() == RECEIVER) continue
                     firTypeBuilder.annotations += annotationEntry.convert<FirAnnotation>()
                 }
             }
+
             return firTypeBuilder.build()
         }
 
diff --git a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt
index 08f30034f39..9f546af1660 100644
--- a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt
+++ b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt
@@ -568,13 +568,11 @@ fun List<FirAnnotationCall>.filterUseSiteTarget(target: AnnotationUseSiteTarget)
         }
     }
 
-fun FirTypeRef.convertToReceiverParameter(): FirReceiverParameter {
+fun FirTypeRef.asReceiverParameter(receiverAnnotations: List<FirAnnotationCall>?): FirReceiverParameter {
     val typeRef = this
     return buildReceiverParameter {
         source = typeRef.source?.fakeElement(KtFakeSourceElementKind.ReceiverFromType)
-        @Suppress("UNCHECKED_CAST")
-        annotations += (typeRef.annotations as List<FirAnnotationCall>).filterUseSiteTarget(AnnotationUseSiteTarget.RECEIVER)
-        (typeRef.annotations as MutableList<FirAnnotation>).removeIf { it.useSiteTarget == AnnotationUseSiteTarget.RECEIVER }
+        receiverAnnotations?.let { annotations += it.filterUseSiteTarget(AnnotationUseSiteTarget.RECEIVER) }
         type = typeRef
     }
 }