Translate attributes in TypeDeserializer

compiler/frontend/src/org/jetbrains/kotlin/frontend/di/injection.kt

@@ -48,6 +48,7 @@ import org.jetbrains.kotlin.types.checker.NewKotlinTypeCheckerImpl
 import org.jetbrains.kotlin.types.expressions.DeclarationScopeProviderForLocalClassifierAnalyzer
 import org.jetbrains.kotlin.types.expressions.LocalClassDescriptorHolder
 import org.jetbrains.kotlin.types.expressions.LocalLazyDeclarationResolver
+import org.jetbrains.kotlin.types.extensions.TypeAttributeTranslators
 import org.jetbrains.kotlin.util.ProgressManagerBasedCancellationChecker
 
 fun StorageComponentContainer.configureModule(
@@ -79,6 +80,8 @@ fun StorageComponentContainer.configureModule(
     analyzerServices.platformConfigurator.configureModuleComponents(this)
     analyzerServices.platformConfigurator.configureModuleDependentCheckers(this)
 
+    useInstance(TypeAttributeTranslators(moduleContext.project).translators)
+
     for (extension in StorageComponentContainerContributor.getInstances(moduleContext.project)) {
         extension.registerModuleComponents(this, platform, moduleContext.module)
     }

core/descriptors/src/org/jetbrains/kotlin/types/TypeAttributeTranslator.kt

@@ -7,12 +7,12 @@ package org.jetbrains.kotlin.types
 
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 
-interface TypeAttributesTranslator {
+interface TypeAttributeTranslator {
     fun toAttributes(annotations: Annotations): TypeAttributes
     fun toAnnotations(attributes: TypeAttributes): Annotations
 }
 
-object DefaultTypeAttributesTranslator : TypeAttributesTranslator {
+object DefaultTypeAttributeTranslator : TypeAttributeTranslator {
     override fun toAnnotations(attributes: TypeAttributes): Annotations {
         return attributes.customAnnotations
     }

core/descriptors/src/org/jetbrains/kotlin/types/TypeAttributes.kt

@@ -114,9 +114,9 @@ class TypeAttributes private constructor(attributes: List<TypeAttribute<*>>) : A
 }
 
 fun TypeAttributes.toDefaultAnnotations(): Annotations =
-    DefaultTypeAttributesTranslator.toAnnotations(this)
+    DefaultTypeAttributeTranslator.toAnnotations(this)
 
-fun Annotations.toDefaultAttributes(): TypeAttributes = DefaultTypeAttributesTranslator.toAttributes(this)
+fun Annotations.toDefaultAttributes(): TypeAttributes = DefaultTypeAttributeTranslator.toAttributes(this)
 
 fun TypeAttributes.replaceAnnotations(newAnnotations: Annotations): TypeAttributes {
     val withoutCustom = (custom?.let { this.remove(it) } ?: this)

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

@@ -534,7 +534,7 @@ interface ClassicTypeSystemContext : TypeSystemInferenceExtensionContext, TypeSy
 
         @Suppress("UNCHECKED_CAST")
         return KotlinTypeFactory.simpleType(
-            DefaultTypeAttributesTranslator.toAttributes(resultingAnnotations),
+            DefaultTypeAttributeTranslator.toAttributes(resultingAnnotations),
             constructor,
             arguments as List<TypeProjection>,
             nullable

core/descriptors/src/org/jetbrains/kotlin/types/extensions/TypeAttributeTranslators.kt

@@ -10,11 +10,12 @@ import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
 import org.jetbrains.kotlin.types.*
 
-interface TypeAttributeTranslatorExtension : TypeAttributesTranslator
+interface TypeAttributeTranslatorExtension : TypeAttributeTranslator
+
+class TypeAttributeTranslators private constructor(val translators: List<TypeAttributeTranslator>) {
+
+    constructor(project: Project) : this(getInstances(project) + DefaultTypeAttributeTranslator)
 
-class TypeAttributeTranslators(project: Project) {
-    val translators: List<TypeAttributesTranslator> =
-        getInstances(project) + DefaultTypeAttributesTranslator
 
     fun toAttributes(annotations: Annotations): TypeAttributes {
         val translated = translators.map { translator ->
@@ -34,5 +35,7 @@ class TypeAttributeTranslators(project: Project) {
         ProjectExtensionDescriptor<TypeAttributeTranslatorExtension>(
             "org.jetbrains.kotlin.extensions.typeAttribute",
             TypeAttributeTranslatorExtension::class.java
-        )
+        ) {
+        val Default = TypeAttributeTranslators(listOf(DefaultTypeAttributeTranslator))
+    }
 }

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/MetadataPackageFragmentProvider.kt

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.serialization.deserialization.builtins.BuiltInSerializerProtocol
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPackageMemberScope
 import org.jetbrains.kotlin.storage.StorageManager
+import org.jetbrains.kotlin.types.TypeAttributeTranslator
 import org.jetbrains.kotlin.types.checker.NewKotlinTypeChecker
 import java.io.InputStream
 
@@ -45,7 +46,8 @@ class MetadataPackageFragmentProvider(
     private val metadataPartProvider: MetadataPartProvider,
     contractDeserializer: ContractDeserializer,
     kotlinTypeChecker: NewKotlinTypeChecker,
-    samConversionResolver: SamConversionResolver
+    samConversionResolver: SamConversionResolver,
+    typeAttributeTranslators: List<TypeAttributeTranslator>
 ) : AbstractDeserializedPackageFragmentProvider(storageManager, finder, moduleDescriptor) {
     init {
         components = DeserializationComponents(
@@ -65,7 +67,8 @@ class MetadataPackageFragmentProvider(
             AdditionalClassPartsProvider.None, PlatformDependentDeclarationFilter.All,
             BuiltInSerializerProtocol.extensionRegistry,
             kotlinTypeChecker,
-            samConversionResolver
+            samConversionResolver,
+            typeAttributeTranslators = typeAttributeTranslators
         )
     }
 

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/TypeDeserializer.kt

@@ -67,6 +67,13 @@ class TypeDeserializer(
         return simpleType(proto, expandTypeAliases = true)
     }
 
+    private fun List<TypeAttributeTranslator>.toAttributes(annotations: Annotations): TypeAttributes {
+        val translated = this.map { translator ->
+            translator.toAttributes(annotations)
+        }.flatten()
+        return TypeAttributes.create(translated)
+    }
+
     fun simpleType(proto: ProtoBuf.Type, expandTypeAliases: Boolean = true): SimpleType {
         val localClassifierType = when {
             proto.hasClassName() -> computeLocalClassifierReplacementType(proto.className)
@@ -85,6 +92,8 @@ class TypeDeserializer(
             c.components.annotationAndConstantLoader.loadTypeAnnotations(proto, c.nameResolver)
         }
 
+        val attributes = c.components.typeAttributeTranslators.toAttributes(annotations)
+
         fun ProtoBuf.Type.collectAllArguments(): List<ProtoBuf.Type.Argument> =
             argumentList + outerType(c.typeTable)?.collectAllArguments().orEmpty()
 
@@ -97,14 +106,17 @@ class TypeDeserializer(
         val simpleType = when {
             expandTypeAliases && declarationDescriptor is TypeAliasDescriptor -> {
                 val expandedType = with(KotlinTypeFactory) { declarationDescriptor.computeExpandedType(arguments) }
+                val expandedAttributes = c.components.typeAttributeTranslators.toAttributes(
+                    Annotations.create(annotations + expandedType.annotations)
+                )
                 expandedType
                     .makeNullableAsSpecified(expandedType.isNullable() || proto.nullable)
-                    .replaceAnnotations(Annotations.create(annotations + expandedType.annotations))
+                    .replaceAttributes(expandedAttributes)
             }
             Flags.SUSPEND_TYPE.get(proto.flags) ->
-                createSuspendFunctionType(annotations, constructor, arguments, proto.nullable)
+                createSuspendFunctionType(attributes, constructor, arguments, proto.nullable)
             else ->
-                KotlinTypeFactory.simpleType(annotations, constructor, arguments, proto.nullable).let {
+                KotlinTypeFactory.simpleType(attributes, constructor, arguments, proto.nullable).let {
                     if (Flags.DEFINITELY_NOT_NULL_TYPE.get(proto.flags))
                         DefinitelyNotNullType.makeDefinitelyNotNull(it) ?: error("null DefinitelyNotNullType for '$it'")
                     else
@@ -157,19 +169,19 @@ class TypeDeserializer(
     }
 
     private fun createSuspendFunctionType(
-        annotations: Annotations,
+        attributes: TypeAttributes,
         functionTypeConstructor: TypeConstructor,
         arguments: List<TypeProjection>,
         nullable: Boolean
     ): SimpleType {
         val result = when (functionTypeConstructor.parameters.size - arguments.size) {
-            0 -> createSuspendFunctionTypeForBasicCase(annotations, functionTypeConstructor, arguments, nullable)
+            0 -> createSuspendFunctionTypeForBasicCase(attributes, functionTypeConstructor, arguments, nullable)
             // This case for types written by eap compiler 1.1
             1 -> {
                 val arity = arguments.size - 1
                 if (arity >= 0) {
                     KotlinTypeFactory.simpleType(
-                        annotations,
+                        attributes,
                         functionTypeConstructor.builtIns.getSuspendFunction(arity).typeConstructor,
                         arguments,
                         nullable
@@ -187,12 +199,12 @@ class TypeDeserializer(
     }
 
     private fun createSuspendFunctionTypeForBasicCase(
-        annotations: Annotations,
+        attributes: TypeAttributes,
         functionTypeConstructor: TypeConstructor,
         arguments: List<TypeProjection>,
         nullable: Boolean
     ): SimpleType? {
-        val functionType = KotlinTypeFactory.simpleType(annotations, functionTypeConstructor, arguments, nullable)
+        val functionType = KotlinTypeFactory.simpleType(attributes, functionTypeConstructor, arguments, nullable)
         return if (!functionType.isFunctionType) null
         else transformRuntimeFunctionTypeToSuspendFunction(functionType)
     }

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/context.kt

@@ -27,11 +27,12 @@ import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite
-import org.jetbrains.kotlin.resolve.SealedClassInheritorsProvider
 import org.jetbrains.kotlin.resolve.constants.ConstantValue
 import org.jetbrains.kotlin.resolve.sam.SamConversionResolver
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 import org.jetbrains.kotlin.storage.StorageManager
+import org.jetbrains.kotlin.types.DefaultTypeAttributeTranslator
+import org.jetbrains.kotlin.types.TypeAttributeTranslator
 import org.jetbrains.kotlin.types.checker.NewKotlinTypeChecker
 
 class DeserializationComponents(
@@ -53,7 +54,8 @@ class DeserializationComponents(
     val extensionRegistryLite: ExtensionRegistryLite,
     val kotlinTypeChecker: NewKotlinTypeChecker = NewKotlinTypeChecker.Default,
     val samConversionResolver: SamConversionResolver,
-    val platformDependentTypeTransformer: PlatformDependentTypeTransformer = PlatformDependentTypeTransformer.None
+    val platformDependentTypeTransformer: PlatformDependentTypeTransformer = PlatformDependentTypeTransformer.None,
+    val typeAttributeTranslators: List<TypeAttributeTranslator> = listOf(DefaultTypeAttributeTranslator)
 ) {
     val classDeserializer: ClassDeserializer = ClassDeserializer(this)