[FIR] Support methods of cone type contexts with annotation markers

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/CompilerConeAttributes.kt

@@ -82,6 +82,17 @@ object CompilerConeAttributes {
 
         override fun toString(): String = "@UnsafeVariance"
     }
+
+    val compilerAttributeByClassId: Map<ClassId, ConeAttribute<*>> = mapOf(
+        Exact.ANNOTATION_CLASS_ID to Exact,
+        NoInfer.ANNOTATION_CLASS_ID to NoInfer,
+        EnhancedNullability.ANNOTATION_CLASS_ID to EnhancedNullability,
+        ExtensionFunctionType.ANNOTATION_CLASS_ID to ExtensionFunctionType,
+        FlexibleNullability.ANNOTATION_CLASS_ID to FlexibleNullability,
+        UnsafeVariance.ANNOTATION_CLASS_ID to UnsafeVariance
+    )
+
+    val compilerAttributeByFqName: Map<FqName, ConeAttribute<*>> = compilerAttributeByClassId.mapKeys { it.key.asSingleFqName() }
 }
 
 val ConeAttributes.exact: CompilerConeAttributes.Exact? by ConeAttributes.attributeAccessor<CompilerConeAttributes.Exact>()

compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/ConeAttributes.kt

@@ -9,11 +9,12 @@ import org.jetbrains.kotlin.fir.utils.AttributeArrayOwner
 import org.jetbrains.kotlin.fir.utils.Protected
 import org.jetbrains.kotlin.fir.utils.TypeRegistry
 import org.jetbrains.kotlin.fir.utils.isEmpty
+import org.jetbrains.kotlin.types.model.AnnotationMarker
 import org.jetbrains.kotlin.utils.addIfNotNull
 import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KClass
 
-abstract class ConeAttribute<T : ConeAttribute<T>> {
+abstract class ConeAttribute<T : ConeAttribute<T>> : AnnotationMarker {
     abstract fun union(other: @UnsafeVariance T?): T?
     abstract fun intersect(other: @UnsafeVariance T?): T?
     abstract fun isSubtypeOf(other: @UnsafeVariance T?): Boolean
@@ -69,6 +70,13 @@ class ConeAttributes private constructor(attributes: List<ConeAttribute<*>>) : A
         return perform(other) { this.intersect(it) }
     }
 
+    fun remove(attribute: ConeAttribute<*>): ConeAttributes {
+        if (arrayMap.isEmpty()) return this
+        val attributes = arrayMap.filter { it != attribute }
+        if (attributes.size == arrayMap.size) return this
+        return create(attributes)
+    }
+
     override fun iterator(): Iterator<ConeAttribute<*>> {
         return arrayMap.iterator()
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceUtils.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.scopes.processOverriddenFunctions
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
+import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.StandardClassIds
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
@@ -45,16 +46,15 @@ fun ConeKotlinType.isKMutableProperty(session: FirSession): Boolean {
 }
 
 private fun ConeKotlinType.functionClassKind(session: FirSession): FunctionClassKind? {
-    val classId = classId(session) ?: return null
-    return FunctionClassKind.byClassNamePrefix(classId.packageFqName, classId.relativeClassName.asString())
+    return classId(session)?.toFunctionClassKind()
+}
+
+private fun ClassId.toFunctionClassKind(): FunctionClassKind? {
+    return FunctionClassKind.byClassNamePrefix(packageFqName, relativeClassName.asString())
 }
 
 fun ConeKotlinType.isBuiltinFunctionalType(session: FirSession): Boolean {
-    val kind = functionClassKind(session) ?: return false
-    return kind == FunctionClassKind.Function ||
-            kind == FunctionClassKind.KFunction ||
-            kind == FunctionClassKind.SuspendFunction ||
-            kind == FunctionClassKind.KSuspendFunction
+    return functionClassKind(session) != null
 }
 
 fun ConeKotlinType.isFunctionalType(session: FirSession): Boolean {
@@ -62,6 +62,10 @@ fun ConeKotlinType.isFunctionalType(session: FirSession): Boolean {
     return kind == FunctionClassKind.Function
 }
 
+fun ConeClassLikeLookupTag.isBuiltinFunctionalType(): Boolean {
+    return classId.toFunctionClassKind() != null
+}
+
 fun ConeKotlinType.isSuspendFunctionType(session: FirSession): Boolean {
     val kind = functionClassKind(session) ?: return false
     return kind == FunctionClassKind.SuspendFunction ||

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.isPrimitiveNumberOrUnsignedNumberType
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.NoSubstitutor
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
+import org.jetbrains.kotlin.fir.resolve.inference.isFunctionalType
 import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
@@ -59,11 +60,23 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
         arguments: List<TypeArgumentMarker>,
         nullable: Boolean,
         isExtensionFunction: Boolean,
-        annotations: List<AnnotationMarker>? // TODO: process annotations
+        annotations: List<AnnotationMarker>?
     ): SimpleTypeMarker {
-        val attributes = if (isExtensionFunction) // TODO: assert correct type constructor
-            ConeAttributes.WithExtensionFunctionType
-        else ConeAttributes.Empty
+        val attributesList = annotations?.filterIsInstanceTo<ConeAttribute<*>, MutableList<ConeAttribute<*>>>(mutableListOf())
+        val attributes: ConeAttributes = if (isExtensionFunction) {
+            require(constructor is ConeClassLikeLookupTag && constructor.isBuiltinFunctionalType())
+            // We don't want to create new instance of ConeAttributes which
+            //   contains only CompilerConeAttributes.ExtensionFunctionType
+            //   to avoid memory consumption
+            if (attributesList != null) {
+                attributesList += CompilerConeAttributes.ExtensionFunctionType
+                ConeAttributes.create(attributesList)
+            } else {
+                ConeAttributes.WithExtensionFunctionType
+            }
+        } else {
+            attributesList?.let { ConeAttributes.create(it) } ?: ConeAttributes.Empty
+        }
         @Suppress("UNCHECKED_CAST")
         return when (constructor) {
             is ConeClassLikeLookupTag -> ConeClassLikeTypeImpl(
@@ -74,7 +87,8 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
             )
             is ConeTypeParameterLookupTag -> ConeTypeParameterTypeImpl(
                 constructor,
-                nullable
+                nullable,
+                attributes
             )
             else -> error("!")
         }
@@ -245,7 +259,8 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
     }
 
     override fun KotlinTypeMarker.removeAnnotations(): KotlinTypeMarker {
-        return this // TODO
+        require(this is ConeKotlinType)
+        return withAttributes(ConeAttributes.Empty)
     }
 
     override fun SimpleTypeMarker.replaceArguments(newArguments: List<TypeArgumentMarker>): SimpleTypeMarker {
@@ -349,8 +364,8 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
     }
 
     override fun KotlinTypeMarker.removeExactAnnotation(): KotlinTypeMarker {
-        // TODO
-        return this
+        require(this is ConeKotlinType)
+        return withAttributes(attributes.remove(CompilerConeAttributes.Exact))
     }
 
     override fun TypeConstructorMarker.toErrorType(): SimpleTypeMarker {
@@ -384,8 +399,7 @@ interface ConeInferenceContext : TypeSystemInferenceExtensionContext, ConeTypeCo
             this,
             {
                 // FIXME supertypes of type constructor contain unsubstituted arguments
-                @Suppress("UNCHECKED_CAST")
-                it.typeConstructor().supertypes() as Collection<ConeKotlinType>
+                it.typeConstructor().supertypes()
             },
             DFS.VisitedWithSet(),
             object : DFS.AbstractNodeHandler<ConeKotlinType, Boolean>() {

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

@@ -10,6 +10,10 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirConstExpression
+import org.jetbrains.kotlin.fir.expressions.FirNamedArgumentExpression
+import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression
+import org.jetbrains.kotlin.fir.expressions.arguments
 import org.jetbrains.kotlin.fir.resolve.correspondingSupertypesCache
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
@@ -349,7 +353,7 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
 
     override fun KotlinTypeMarker.getAnnotations(): List<AnnotationMarker> {
         require(this is ConeKotlinType)
-        return emptyList() // TODO
+        return attributes.arrayMap.toList()
     }
 
     override fun SimpleTypeMarker.isStubType(): Boolean {
@@ -425,13 +429,30 @@ interface ConeTypeContext : TypeSystemContext, TypeSystemOptimizationContext, Ty
     }
 
     override fun KotlinTypeMarker.hasAnnotation(fqName: FqName): Boolean {
-        // TODO support annotations
-        return false
+        require(this is ConeKotlinType)
+        val compilerAttribute = CompilerConeAttributes.compilerAttributeByFqName[fqName]
+        if (compilerAttribute != null) {
+            return compilerAttribute in attributes
+        }
+        val customAnnotations = attributes.customAnnotations
+        return customAnnotations.any {
+            it.typeRef.coneTypeSafe<ConeKotlinType>()?.fullyExpandedType(session)?.classId?.asSingleFqName() == fqName
+        }
     }
 
     override fun KotlinTypeMarker.getAnnotationFirstArgumentValue(fqName: FqName): Any? {
-        // TODO support annotations
-        return null
+        require(this is ConeKotlinType)
+        // We don't check for compiler attributes because all of them doesn't have parameters
+        val customAnnotations = attributes.customAnnotations
+        val annotationCall = customAnnotations.firstOrNull {
+            it.typeRef.coneTypeSafe<ConeKotlinType>()?.fullyExpandedType(session)?.classId?.asSingleFqName() == fqName
+        } ?: return null
+        val argument = when (val argument = annotationCall.arguments.firstOrNull() ?: return null) {
+            is FirVarargArgumentsExpression -> argument.arguments.firstOrNull()
+            is FirNamedArgumentExpression -> argument.expression
+            else -> argument
+        } ?: return null
+        return (argument as? FirConstExpression<*>)?.value
     }
 
     override fun TypeConstructorMarker.getTypeParameterClassifier(): TypeParameterMarker? {