[FIR] Support typealiases in FirTypeAnnotationChecker

`directExpansionType` is needed here for cases like
`typealias S = @S Suppress`. Otherwise, `fullyExpandedType`
would cache the same typealias

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/type/FirTypeAnnotationChecker.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.fir.analysis.checkers.getAllowedAnnotationTargets
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.diagnostics.reportOn
-import org.jetbrains.kotlin.fir.expressions.classId
+import org.jetbrains.kotlin.fir.declarations.toAnnotationClassId
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.isBuiltinFunctionalType
@@ -32,7 +32,7 @@ object FirTypeAnnotationChecker : FirTypeRefChecker() {
                     reporter.reportOn(annotation.source, FirErrors.WRONG_ANNOTATION_TARGET, "type usage", context)
                 }
             }
-            if (annotation.classId == StandardClassIds.Annotations.ExtensionFunctionType) {
+            if (annotation.toAnnotationClassId(context.session) == StandardClassIds.Annotations.ExtensionFunctionType) {
                 if (!typeRef.type.isBuiltinFunctionalType(context.session)) {
                     if (context.languageVersionSettings.supportsFeature(LanguageFeature.ForbidExtensionFunctionTypeOnNonFunctionTypes)) {
                         reporter.reportOn(annotation.source, FirErrors.WRONG_EXTENSION_FUNCTION_TYPE, context)

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/FirTypeDeserializer.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.deserialization
 
 import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
 import org.jetbrains.kotlin.fir.FirModuleData
+import org.jetbrains.kotlin.fir.computeTypeAttributes
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.FirTypeParameterRefsOwner

compiler/fir/providers/src/org/jetbrains/kotlin/fir/CopyUtils.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir
 
+import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.descriptors.EffectiveVisibility
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibility
@@ -14,13 +15,19 @@ import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
 import org.jetbrains.kotlin.fir.declarations.builder.buildTypeParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.declarations.utils.expandedConeType
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
+import org.jetbrains.kotlin.fir.extensions.extensionService
+import org.jetbrains.kotlin.fir.extensions.typeAttributeExtensions
 import org.jetbrains.kotlin.fir.references.FirReference
+import org.jetbrains.kotlin.fir.resolve.directExpansionType
+import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.name.ClassId
 
 inline fun FirFunctionCall.copyAsImplicitInvokeCall(
     setupCopy: FirImplicitInvokeCallBuilder.() -> Unit
@@ -72,3 +79,51 @@ fun FirTypeRef.errorTypeFromPrototype(
         this.diagnostic = diagnostic
     }
 }
+
+fun List<FirAnnotation>.computeTypeAttributes(session: FirSession, predefined: List<ConeAttribute<*>> = emptyList()): ConeAttributes {
+    if (this.isEmpty()) {
+        if (predefined.isEmpty()) return ConeAttributes.Empty
+        return ConeAttributes.create(predefined)
+    }
+    val attributes = mutableListOf<ConeAttribute<*>>()
+    attributes += predefined
+    val customAnnotations = mutableListOf<FirAnnotation>()
+    for (annotation in this) {
+        when (annotation.tryExpandClassId(session)) {
+            CompilerConeAttributes.Exact.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.Exact
+            CompilerConeAttributes.NoInfer.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.NoInfer
+            CompilerConeAttributes.ExtensionFunctionType.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.ExtensionFunctionType
+            CompilerConeAttributes.ContextFunctionTypeParams.ANNOTATION_CLASS_ID ->
+                attributes +=
+                    CompilerConeAttributes.ContextFunctionTypeParams(
+                        annotation.extractContextReceiversCount() ?: 0
+                    )
+
+            CompilerConeAttributes.UnsafeVariance.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.UnsafeVariance
+            else -> {
+                val attributeFromPlugin = session.extensionService.typeAttributeExtensions.firstNotNullOfOrNull {
+                    it.extractAttributeFromAnnotation(annotation)
+                }
+                if (attributeFromPlugin != null) {
+                    attributes += attributeFromPlugin
+                } else {
+                    customAnnotations += annotation
+                }
+            }
+        }
+    }
+    if (customAnnotations.isNotEmpty()) {
+        attributes += CustomAnnotationTypeAttribute(customAnnotations)
+    }
+    return ConeAttributes.create(attributes)
+}
+
+private fun FirAnnotation.tryExpandClassId(session: FirSession): ClassId? {
+    return when (val directlyExpanded = coneClassLikeType?.directExpansionType(session) { it.expandedConeType }) {
+        null -> coneClassLikeType?.classId // mutually recursive typealiases
+        else -> directlyExpanded.fullyExpandedType(session).classId
+    }
+}
+
+private fun FirAnnotation.extractContextReceiversCount() =
+    (argumentMapping.mapping[StandardNames.CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME] as? FirConstExpression<*>)?.value as? Int

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirTypeUtils.kt

@@ -5,11 +5,7 @@
 
 package org.jetbrains.kotlin.fir.types
 
-import org.jetbrains.kotlin.builtins.StandardNames
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.typeAttributeExtensions
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
 import org.jetbrains.kotlin.name.ClassId
@@ -112,47 +108,6 @@ fun ConeClassLikeType.toConstKind(): ConstantValueKind<*>? = when (lookupTag.cla
     else -> null
 }
 
-fun List<FirAnnotation>.computeTypeAttributes(session: FirSession, predefined: List<ConeAttribute<*>> = emptyList()): ConeAttributes {
-    if (this.isEmpty()) {
-        if (predefined.isEmpty()) return ConeAttributes.Empty
-        return ConeAttributes.create(predefined)
-    }
-    val attributes = mutableListOf<ConeAttribute<*>>()
-    attributes += predefined
-    val customAnnotations = mutableListOf<FirAnnotation>()
-    for (annotation in this) {
-        val type = annotation.annotationTypeRef.coneTypeSafe<ConeClassLikeType>() ?: continue
-        when (type.lookupTag.classId) {
-            CompilerConeAttributes.Exact.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.Exact
-            CompilerConeAttributes.NoInfer.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.NoInfer
-            CompilerConeAttributes.ExtensionFunctionType.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.ExtensionFunctionType
-            CompilerConeAttributes.ContextFunctionTypeParams.ANNOTATION_CLASS_ID ->
-                attributes +=
-                    CompilerConeAttributes.ContextFunctionTypeParams(
-                        annotation.extractContextReceiversCount() ?: 0
-                    )
-            CompilerConeAttributes.UnsafeVariance.ANNOTATION_CLASS_ID -> attributes += CompilerConeAttributes.UnsafeVariance
-            else -> {
-                val attributeFromPlugin = session.extensionService.typeAttributeExtensions.firstNotNullOfOrNull {
-                    it.extractAttributeFromAnnotation(annotation)
-                }
-                if (attributeFromPlugin != null) {
-                    attributes += attributeFromPlugin
-                } else {
-                    customAnnotations += annotation
-                }
-            }
-        }
-    }
-    if (customAnnotations.isNotEmpty()) {
-        attributes += CustomAnnotationTypeAttribute(customAnnotations)
-    }
-    return ConeAttributes.create(attributes)
-}
-
-private fun FirAnnotation.extractContextReceiversCount() =
-    (argumentMapping.mapping[StandardNames.CONTEXT_FUNCTION_TYPE_PARAMETER_COUNT_NAME] as? FirConstExpression<*>)?.value as? Int
-
 fun FirTypeProjection.toConeTypeProjection(): ConeTypeProjection =
     when (this) {
         is FirStarProjection -> ConeStarProjection

compiler/testData/diagnostics/tests/FreeFunctionCalledAsExtension.fir.kt

@@ -5,8 +5,9 @@ fun foo(a: (String) -> Unit) {
 
 
 interface A : (String) -> Unit {}
+typealias AliasedEFT = ExtensionFunctionType
 
-fun foo(a: <!WRONG_EXTENSION_FUNCTION_TYPE_WARNING!>@ExtensionFunctionType<!> A) {
+fun foo(a: <!WRONG_EXTENSION_FUNCTION_TYPE_WARNING!>@AliasedEFT<!> A) {
     // @Extension annotation on an unrelated type shouldn't have any effect on this diagnostic.
     // Only kotlin.Function{n} type annotated with @Extension should
     "".a()

compiler/testData/diagnostics/tests/FreeFunctionCalledAsExtension.kt

@@ -5,8 +5,9 @@ fun foo(a: (String) -> Unit) {
 
 
 interface A : (String) -> Unit {}
+typealias AliasedEFT = ExtensionFunctionType
 
-fun foo(a: <!WRONG_EXTENSION_FUNCTION_TYPE_WARNING!>@ExtensionFunctionType<!> A) {
+fun foo(a: <!WRONG_EXTENSION_FUNCTION_TYPE_WARNING!>@AliasedEFT<!> A) {
     // @Extension annotation on an unrelated type shouldn't have any effect on this diagnostic.
     // Only kotlin.Function{n} type annotated with @Extension should
     "".<!UNRESOLVED_REFERENCE!>a<!>()

compiler/testData/diagnostics/tests/FreeFunctionCalledAsExtension.txt

@@ -1,7 +1,7 @@
 package
 
 public fun foo(/*0*/ a: (kotlin.String) -> kotlin.Unit): kotlin.Unit
-public fun foo(/*0*/ a: @kotlin.ExtensionFunctionType A): kotlin.Unit
+public fun foo(/*0*/ a: @AliasedEFT /* = kotlin.ExtensionFunctionType */ A): kotlin.Unit
 
 public interface A : (kotlin.String) -> kotlin.Unit {
     public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
@@ -9,3 +9,5 @@ public interface A : (kotlin.String) -> kotlin.Unit {
     public abstract override /*1*/ /*fake_override*/ fun invoke(/*0*/ p1: kotlin.String): kotlin.Unit
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
 }
+public typealias AliasedEFT = kotlin.ExtensionFunctionType
+