Propagate all annotations during creating simple functional types

^KT-44563 Fixed

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeInferenceContext.kt

@@ -58,7 +58,8 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
         constructor: TypeConstructorMarker,
         arguments: List<TypeArgumentMarker>,
         nullable: Boolean,
-        isExtensionFunction: Boolean
+        isExtensionFunction: Boolean,
+        annotations: List<AnnotationMarker>? // TODO: process annotations
     ): SimpleTypeMarker {
         val attributes = if (isExtensionFunction) // TODO: assert correct type constructor
             ConeAttributes.create(listOf(CompilerConeAttributes.ExtensionFunctionType))

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/ConeTypeContext.kt

@@ -402,6 +402,11 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
         return false
     }
 
+    override fun KotlinTypeMarker.getAnnotations(): List<AnnotationMarker> {
+        require(this is ConeKotlinType)
+        return emptyList() // TODO
+    }
+
     override fun SimpleTypeMarker.isStubType(): Boolean {
         return this is StubTypeMarker
     }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/types/IrTypeSystemContext.kt

@@ -10,9 +10,12 @@ import org.jetbrains.kotlin.builtins.PrimitiveType
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.expressions.IrConst
+import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
 import org.jetbrains.kotlin.ir.symbols.FqNameEqualityChecker
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
@@ -222,13 +225,22 @@ interface IrTypeSystemContext : TypeSystemContext, TypeSystemCommonSuperTypesCon
         return IrDynamicTypeImpl(null, emptyList(), Variance.INVARIANT)
     }
 
-    // TODO: implement taking into account `isExtensionFunction`
     override fun createSimpleType(
         constructor: TypeConstructorMarker,
         arguments: List<TypeArgumentMarker>,
         nullable: Boolean,
-        isExtensionFunction: Boolean
-    ): SimpleTypeMarker = IrSimpleTypeImpl(constructor as IrClassifierSymbol, nullable, arguments.map { it as IrTypeArgument }, emptyList())
+        isExtensionFunction: Boolean,
+        annotations: List<AnnotationMarker>?
+    ): SimpleTypeMarker {
+        val ourAnnotations = annotations?.filterIsInstance<IrConstructorCall>()
+        require(ourAnnotations?.size == annotations?.size)
+        return IrSimpleTypeImpl(
+            constructor as IrClassifierSymbol,
+            nullable,
+            arguments.map { it as IrTypeArgument },
+            ourAnnotations ?: emptyList()
+        )
+    }
 
     private fun TypeVariance.convertVariance(): Variance {
         return when (this) {
@@ -283,6 +295,11 @@ interface IrTypeSystemContext : TypeSystemContext, TypeSystemCommonSuperTypesCon
     override fun SimpleTypeMarker.isPrimitiveType(): Boolean =
         this is IrSimpleType && irTypePredicates_isPrimitiveType()
 
+    override fun KotlinTypeMarker.getAnnotations(): List<AnnotationMarker> {
+        require(this is IrType)
+        return this.annotations.map { object : AnnotationMarker, IrElement by it {} }
+    }
+
     override fun createErrorType(debugName: String): SimpleTypeMarker {
         TODO("IrTypeSystemContext doesn't support constraint system resolution")
     }

compiler/resolution.common/src/org/jetbrains/kotlin/resolve/calls/inference/components/PostponedArgumentInputTypesResolver.kt

@@ -24,6 +24,7 @@ class PostponedArgumentInputTypesResolver(
         val parametersFromDeclaration: List<KotlinTypeMarker?>?,
         val parametersFromDeclarationOfRelatedLambdas: Set<List<KotlinTypeMarker?>>?,
         val parametersFromConstraints: Set<List<TypeWithKind>>?,
+        val annotations: List<AnnotationMarker>?,
         val isExtensionFunction: Boolean,
         val isSuspend: Boolean,
         val isNullable: Boolean
@@ -57,7 +58,7 @@ class PostponedArgumentInputTypesResolver(
         argument: PostponedAtomWithRevisableExpectedType,
         postponedArguments: List<PostponedAtomWithRevisableExpectedType>,
         variableDependencyProvider: TypeVariableDependencyInformationProvider
-    ): ParameterTypesInfo? = with(resolutionTypeSystemContext) {
+    ): ParameterTypesInfo? {
         val expectedType = argument.expectedType ?: return null
         val variableWithConstraints = notFixedTypeVariables[expectedType.typeConstructor()] ?: return null
         val functionalTypesFromConstraints = findFunctionalTypesInConstraints(variableWithConstraints, variableDependencyProvider)
@@ -76,6 +77,8 @@ class PostponedArgumentInputTypesResolver(
             }
         }
 
+        val annotations = functionalTypesFromConstraints?.map { it.type.getAnnotations() }?.flatten()?.distinct()
+
         // An extension function flag can only come from a declaration of anonymous function: `select({ this + it }, fun Int.(x: Int) = 10)`
         val (parameterTypesFromDeclarationOfRelatedLambdas, isThereExtensionFunctionAmongRelatedLambdas) =
             getDeclaredParametersFromRelatedLambdas(argument, postponedArguments, variableDependencyProvider)
@@ -96,6 +99,7 @@ class PostponedArgumentInputTypesResolver(
             parameterTypesFromDeclaration,
             parameterTypesFromDeclarationOfRelatedLambdas,
             parameterTypesFromConstraints,
+            annotations = annotations,
             isExtensionFunction = isThereExtensionFunctionAmongRelatedLambdas || extensionFunctionTypePresentInConstraints,
             isSuspend = isSuspend,
             isNullable = isNullable
@@ -344,7 +348,8 @@ class PostponedArgumentInputTypesResolver(
                 shouldDiscriminateExtensionFunctionAnnotation -> false
                 argument.isFunctionExpressionWithReceiver() -> true
                 else -> parameterTypesInfo.isExtensionFunction
-            }
+            },
+            annotations = parameterTypesInfo.annotations
         )
 
         getBuilder().addSubtypeConstraint(

compiler/testData/diagnostics/tests/annotations/functionalTypes/propagteAnyAnnotations.fir.kt

@@ -0,0 +1,12 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.TYPE)
+annotation class Composable
+
+fun bar(p: @Composable ()->Unit) {}
+
+@Composable fun foo() {}
+
+fun main() {
+    bar(if(true) { { foo() } } else { { } })
+}

compiler/testData/diagnostics/tests/annotations/functionalTypes/propagteAnyAnnotations.kt

@@ -0,0 +1,12 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.TYPE)
+annotation class Composable
+
+fun bar(p: @Composable ()->Unit) {}
+
+@Composable fun foo() {}
+
+fun main() {
+    bar(<!DEBUG_INFO_EXPRESSION_TYPE("@Composable() () -> kotlin.Unit")!>if(true) { { foo() } } else { { } }<!>)
+}

compiler/testData/diagnostics/tests/annotations/functionalTypes/propagteAnyAnnotations.txt

@@ -0,0 +1,12 @@
+package
+
+public fun bar(/*0*/ p: @Composable() () -> kotlin.Unit): kotlin.Unit
+@Composable public fun foo(): kotlin.Unit
+public fun main(): kotlin.Unit
+
+@kotlin.annotation.Target(allowedTargets = {AnnotationTarget.FUNCTION, AnnotationTarget.TYPE}) public final annotation class Composable : kotlin.Annotation {
+    public constructor Composable()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

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)
     }