FIR: Serialize type attributes for non-type-ref types, too

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirElementSerializer.kt

@@ -622,39 +622,7 @@ class FirElementSerializer private constructor(
     fun typeId(type: ConeKotlinType): Int = typeTable[typeProto(type)]
 
     private fun typeProto(typeRef: FirTypeRef, toSuper: Boolean = false): ProtoBuf.Type.Builder {
-        val coneType = typeRef.coneType
+        return typeProto(typeRef.coneType, toSuper, correspondingTypeRef = typeRef)
-        return typeProto(coneType, toSuper, correspondingTypeRef = typeRef).also { typeProto ->
-            val compilerAttributes = mutableListOf<ConeAttribute<*>>()
-            val extensionAttributes = mutableListOf<ConeAttribute<*>>()
-            for (attribute in coneType.attributes) {
-                when {
-                    attribute is CustomAnnotationTypeAttribute -> continue
-                    attribute.key in CompilerConeAttributes.classIdByCompilerAttributeKey -> compilerAttributes += attribute
-                    else -> extensionAttributes += attribute
-                }
-            }
-
-            for (attribute in compilerAttributes) {
-                val annotation = buildAnnotation {
-                    annotationTypeRef = buildResolvedTypeRef {
-                        type = ConeClassLikeTypeImpl(
-                            ConeClassLikeLookupTagImpl(CompilerConeAttributes.classIdByCompilerAttributeKey.getValue(attribute.key)),
-                            emptyArray(),
-                            isNullable = false
-                        )
-                    }
-                    argumentMapping = FirEmptyAnnotationArgumentMapping
-                }
-                extension.serializeTypeAnnotation(annotation, typeProto)
-            }
-
-            for (attributeExtension in session.extensionService.typeAttributeExtensions) {
-                for (attribute in extensionAttributes) {
-                    val annotation = attributeExtension.convertAttributeToAnnotation(attribute) ?: continue
-                    extension.serializeTypeAnnotation(annotation, typeProto)
-                }
-            }
-        }
     }
 
     private fun typeProto(
@@ -739,8 +707,24 @@ class FirElementSerializer private constructor(
             builder.nullable = type.isMarkedNullable
         }
 
-        for (annotation in type.attributes.customAnnotations) {
+        val extensionAttributes = mutableListOf<ConeAttribute<*>>()
-            extension.serializeTypeAnnotation(annotation, builder)
+        for (attribute in type.attributes) {
+            when {
+                attribute is CustomAnnotationTypeAttribute ->
+                    for (annotation in attribute.annotations) {
+                        extension.serializeTypeAnnotation(annotation, builder)
+                    }
+                attribute.key in CompilerConeAttributes.classIdByCompilerAttributeKey ->
+                    serializeCompilerDefinedTypeAttribute(builder, attribute)
+                else -> extensionAttributes += attribute
+            }
+        }
+
+        for (attributeExtension in session.extensionService.typeAttributeExtensions) {
+            for (attribute in extensionAttributes) {
+                val annotation = attributeExtension.convertAttributeToAnnotation(attribute) ?: continue
+                extension.serializeTypeAnnotation(annotation, builder)
+            }
         }
 
         // TODO: abbreviated type
@@ -756,6 +740,23 @@ class FirElementSerializer private constructor(
         return builder
     }
 
+    private fun serializeCompilerDefinedTypeAttribute(
+        builder: ProtoBuf.Type.Builder,
+        attribute: ConeAttribute<*>
+    ) {
+        val annotation = buildAnnotation {
+            annotationTypeRef = buildResolvedTypeRef {
+                this.type = ConeClassLikeTypeImpl(
+                    ConeClassLikeLookupTagImpl(CompilerConeAttributes.classIdByCompilerAttributeKey.getValue(attribute.key)),
+                    emptyArray(),
+                    isNullable = false
+                )
+            }
+            argumentMapping = FirEmptyAnnotationArgumentMapping
+        }
+        extension.serializeTypeAnnotation(annotation, builder)
+    }
+
     private fun serializeAnnotationFromAttribute(
         existingAnnotations: List<FirAnnotation>?,
         classId: ClassId,