Propagate all annotations during creating simple functional types

^KT-44563 Fixed

core/compiler.common/src/org/jetbrains/kotlin/types/model/TypeSystemContext.kt

@@ -36,6 +36,7 @@ interface IntersectionTypeConstructorMarker : TypeConstructorMarker
 
 interface TypeSubstitutorMarker
 
+interface AnnotationMarker
 
 enum class TypeVariance(val presentation: String) {
     IN("in"),
@@ -73,7 +74,8 @@ interface TypeSystemTypeFactoryContext {
         constructor: TypeConstructorMarker,
         arguments: List<TypeArgumentMarker>,
         nullable: Boolean,
-        isExtensionFunction: Boolean = false
+        isExtensionFunction: Boolean = false,
+        annotations: List<AnnotationMarker>? = null
     ): SimpleTypeMarker
 
     fun createTypeArgument(type: KotlinTypeMarker, variance: TypeVariance): TypeArgumentMarker
@@ -391,6 +393,8 @@ interface TypeSystemContext : TypeSystemOptimizationContext {
     fun prepareType(type: KotlinTypeMarker): KotlinTypeMarker
 
     fun SimpleTypeMarker.isPrimitiveType(): Boolean
+
+    fun KotlinTypeMarker.getAnnotations(): List<AnnotationMarker>
 }
 
 enum class CaptureStatus {

core/descriptors/src/org/jetbrains/kotlin/descriptors/annotations/AnnotationDescriptor.kt

@@ -25,8 +25,9 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameOrNull
 import org.jetbrains.kotlin.types.ErrorUtils
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.getAbbreviation
+import org.jetbrains.kotlin.types.model.AnnotationMarker
 
-interface AnnotationDescriptor {
+interface AnnotationDescriptor : AnnotationMarker {
     val type: KotlinType
 
     val fqName: FqName?

core/descriptors/src/org/jetbrains/kotlin/descriptors/annotations/Annotations.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.descriptors.annotations
 
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.types.model.AnnotationMarker
 
 interface Annotated {
     val annotations: Annotations

core/descriptors/src/org/jetbrains/kotlin/types/checker/ClassicTypeSystemContext.kt

@@ -29,8 +29,6 @@ import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.types.typeUtil.contains
 import org.jetbrains.kotlin.types.typeUtil.representativeUpperBound
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
-import kotlin.contracts.ExperimentalContracts
-import kotlin.contracts.contract
 import org.jetbrains.kotlin.types.typeUtil.isSignedOrUnsignedNumberType as classicIsSignedOrUnsignedNumberType
 
 interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSystemCommonBackendContext {
@@ -437,16 +435,25 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
         constructor: TypeConstructorMarker,
         arguments: List<TypeArgumentMarker>,
         nullable: Boolean,
-        isExtensionFunction: Boolean
+        isExtensionFunction: Boolean,
+        annotations: List<AnnotationMarker>?
     ): SimpleTypeMarker {
         require(constructor is TypeConstructor, constructor::errorMessage)
 
-        val annotations = if (isExtensionFunction) {
-            Annotations.create(listOf(BuiltInAnnotationDescriptor(builtIns, FqNames.extensionFunctionType, emptyMap())))
-        } else Annotations.EMPTY
+        val ourAnnotations = annotations?.filterIsInstance<AnnotationDescriptor>()
+        require(ourAnnotations?.size == annotations?.size)
+
+        fun createExtensionFunctionAnnotation() = BuiltInAnnotationDescriptor(builtIns, FqNames.extensionFunctionType, emptyMap())
+
+        val resultingAnnotations = when {
+            ourAnnotations.isNullOrEmpty() && isExtensionFunction -> Annotations.create(listOf(createExtensionFunctionAnnotation()))
+            !ourAnnotations.isNullOrEmpty() && !isExtensionFunction -> Annotations.create(ourAnnotations.filter { it.fqName != FqNames.extensionFunctionType })
+            !ourAnnotations.isNullOrEmpty() && isExtensionFunction -> Annotations.create(ourAnnotations + createExtensionFunctionAnnotation())
+            else -> Annotations.EMPTY
+        }
 
         @Suppress("UNCHECKED_CAST")
-        return KotlinTypeFactory.simpleType(annotations, constructor, arguments as List<TypeProjection>, nullable)
+        return KotlinTypeFactory.simpleType(resultingAnnotations, constructor, arguments as List<TypeProjection>, nullable)
     }
 
     override fun createTypeArgument(type: KotlinTypeMarker, variance: TypeVariance): TypeArgumentMarker {
@@ -557,6 +564,11 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
         return KotlinBuiltIns.isPrimitiveType(this)
     }
 
+    override fun KotlinTypeMarker.getAnnotations(): List<AnnotationMarker> {
+        require(this is KotlinType, this::errorMessage)
+        return this.annotations.toList()
+    }
+
     override fun captureFromExpression(type: KotlinTypeMarker): KotlinTypeMarker? {
         return captureFromExpressionInternal(type as UnwrappedType)
     }