Klib metadata cached package fragment and package fragment provider

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/LegacyDescriptorUtils.kt

@@ -6,17 +6,12 @@
 package org.jetbrains.kotlin.backend.common.serialization
 
 import org.jetbrains.kotlin.backend.common.descriptors.propertyIfAccessor
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.MemberDescriptor
-import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
-import org.jetbrains.kotlin.resolve.checkers.ExpectedActualDeclarationChecker
 import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.library.KotlinLibrary
+import org.jetbrains.kotlin.library.metadata.KlibMetadataDeserializedPackageFragment
-import org.jetbrains.kotlin.library.metadata.KlibMetadataPackageFragment
 import org.jetbrains.kotlin.library.metadata.KlibMetadataProtoBuf
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.OverridingUtil
+import org.jetbrains.kotlin.resolve.checkers.ExpectedActualDeclarationChecker
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPropertyDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedSimpleFunctionDescriptor
 
@@ -56,7 +51,7 @@ val ModuleDescriptor.isForwardDeclarationModule get() =
     name == Name.special("<forward declarations>")
 
 private fun sourceByIndex(descriptor: CallableMemberDescriptor, index: Int): SourceFile {
-    val fragment = descriptor.findPackage() as KlibMetadataPackageFragment
+    val fragment = descriptor.findPackage() as KlibMetadataDeserializedPackageFragment
     return fragment.fileRegistry.sourceFile(index)
 }
 

compiler/util-klib-metadata/src/org/jetbrains/kotlin/library/metadata/KlibMetadataDeserializationUtils.kt

@@ -22,10 +22,18 @@ interface PackageAccessHandler {
         library: KotlinLibrary,
         packageFqName: String,
         partName: String
-    ): ProtoBuf.PackageFragment = parsePackageFragment(library.packageMetadata(packageFqName, partName))
+    ): ProtoBuf.PackageFragment = loadPackageFragmentByteArray(library.packageMetadata(packageFqName, partName))
+
+    fun loadPackageFragmentByteArray(byteArray: ByteArray): ProtoBuf.PackageFragment = parsePackageFragment(byteArray)
 
     fun markNeededForLink(library: KotlinLibrary, fqName: String) {}
 }
 
-object SimplePackageAccessHandler : PackageAccessHandler
+object SimplePackageAccessHandler : PackageAccessHandler {
+    override fun loadPackageFragment(
+        library: KotlinLibrary,
+        packageFqName: String,
+        partName: String
+    ): ProtoBuf.PackageFragment = parsePackageFragment(library.packageMetadata(packageFqName, partName))
+}
 

compiler/util-klib-metadata/src/org/jetbrains/kotlin/library/metadata/KlibMetadataDeserializedPackageFragmentsFactory.kt

@@ -16,6 +16,12 @@ interface KlibMetadataDeserializedPackageFragmentsFactory {
         storageManager: StorageManager
     ): List<KlibMetadataPackageFragment>
 
+    fun createCachedPackageFragments(
+        packageFragments: List<ByteArray>,
+        moduleDescriptor: ModuleDescriptor,
+        storageManager: StorageManager
+    ): List<KlibMetadataPackageFragment>
+
     fun createSyntheticPackageFragments(
         library: KotlinLibrary,
         deserializedPackageFragments: List<KlibMetadataPackageFragment>,

compiler/util-klib-metadata/src/org/jetbrains/kotlin/library/metadata/KlibMetadataModuleDescriptorFactory.kt

@@ -56,4 +56,11 @@ interface KlibMetadataModuleDescriptorFactory {
         configuration: DeserializationConfiguration,
         compositePackageFragmentAddend: PackageFragmentProvider?
     ): PackageFragmentProvider
+
+    fun createCachedPackageFragmentProvider(
+        byteArrays: List<ByteArray>,
+        storageManager: StorageManager,
+        moduleDescriptor: ModuleDescriptor,
+        configuration: DeserializationConfiguration
+    ): PackageFragmentProvider
 }

compiler/util-klib-metadata/src/org/jetbrains/kotlin/library/metadata/KlibMetadataPackageFragment.kt

@@ -29,19 +29,47 @@ private val KotlinLibrary.fileSources: KlibMetadataFileRegistry get() {
 }
 
 
-class KlibMetadataPackageFragment(
+class KlibMetadataDeserializedPackageFragment(
     fqName: FqName,
     private val library: KotlinLibrary,
     private val packageAccessHandler: PackageAccessHandler?,
     storageManager: StorageManager,
     module: ModuleDescriptor,
     private val partName: String
-) : DeserializedPackageFragment(fqName, storageManager, module) {
+) : KlibMetadataPackageFragment(fqName, storageManager, module) {
+
+    // The proto field is lazy so that we can load only needed
+    // packages from the library.
+    override val protoForNames: ProtoBuf.PackageFragment by lazy {
+        (packageAccessHandler ?: SimplePackageAccessHandler).loadPackageFragment(library, fqName.asString(), partName)
+    }
 
     val fileRegistry: KlibMetadataFileRegistry by lazy {
-            library.fileSources
+        library.fileSources
     }
 
+
+    override val proto: ProtoBuf.PackageFragment
+        get() {
+            packageAccessHandler?.markNeededForLink(library, fqName.asString())
+            return protoForNames
+        }
+}
+
+class KlibMetadataCachedPackageFragment(
+    byteArray: ByteArray,
+    storageManager: StorageManager,
+    module: ModuleDescriptor,
+    override val protoForNames: ProtoBuf.PackageFragment = parsePackageFragment(byteArray),
+    fqName: FqName = FqName(protoForNames.getExtension(KlibMetadataProtoBuf.fqName))
+) :  KlibMetadataPackageFragment(fqName, storageManager, module)
+
+abstract class KlibMetadataPackageFragment(
+    fqName: FqName,
+    storageManager: StorageManager,
+    module: ModuleDescriptor
+) : DeserializedPackageFragment(fqName, storageManager, module) {
+
     lateinit var components: DeserializationComponents
 
     override fun initialize(components: DeserializationComponents) {
@@ -50,15 +78,10 @@ class KlibMetadataPackageFragment(
 
     // The proto field is lazy so that we can load only needed
     // packages from the library.
-    private val protoForNames: ProtoBuf.PackageFragment by lazy {
+    abstract val protoForNames: ProtoBuf.PackageFragment
-        (packageAccessHandler ?: SimplePackageAccessHandler).loadPackageFragment(library, fqName.asString(), partName)
-    }
 
-    val proto: ProtoBuf.PackageFragment
+    open val proto: ProtoBuf.PackageFragment
-        get() {
+        get() = protoForNames
-            packageAccessHandler?.markNeededForLink(library, fqName.asString())
-            return protoForNames
-        }
 
     private val nameResolver by lazy {
         NameResolverImpl(protoForNames.strings, protoForNames.qualifiedNames)

compiler/util-klib-metadata/src/org/jetbrains/kotlin/library/metadata/impl/KlibMetadataDeserializedPackageFragmentsFactoryImpl.kt

@@ -7,6 +7,8 @@ import org.jetbrains.kotlin.descriptors.impl.PackageFragmentDescriptorImpl
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.library.*
+import org.jetbrains.kotlin.library.metadata.KlibMetadataCachedPackageFragment
+import org.jetbrains.kotlin.library.metadata.KlibMetadataDeserializedPackageFragment
 import org.jetbrains.kotlin.library.metadata.KlibMetadataPackageFragment
 import org.jetbrains.kotlin.library.metadata.PackageAccessHandler
 import org.jetbrains.kotlin.name.FqName
@@ -34,10 +36,17 @@ open class KlibMetadataDeserializedPackageFragmentsFactoryImpl : KlibMetadataDes
         val fqName = FqName(it)
         val parts = library.packageMetadataParts(fqName.asString())
         parts.map { partName ->
-            KlibMetadataPackageFragment(fqName, library, packageAccessedHandler, storageManager, moduleDescriptor, partName)
+            KlibMetadataDeserializedPackageFragment(fqName, library, packageAccessedHandler, storageManager, moduleDescriptor, partName)
         }
     }
 
+    override fun createCachedPackageFragments(
+        packageFragments: List<ByteArray>,
+        moduleDescriptor: ModuleDescriptor,
+        storageManager: StorageManager
+    ) = packageFragments.map { byteArray ->
+        KlibMetadataCachedPackageFragment(byteArray, storageManager, moduleDescriptor)
+    }
 
     override fun createSyntheticPackageFragments(
         library: KotlinLibrary,

compiler/util-klib-metadata/src/org/jetbrains/kotlin/library/metadata/impl/KlibMetadataModuleDescriptorFactoryImpl.kt

@@ -78,6 +78,22 @@ class KlibMetadataModuleDescriptorFactoryImpl(
         return moduleDescriptor
     }
 
+    override fun createCachedPackageFragmentProvider(
+        byteArrays: List<ByteArray>,
+        storageManager: StorageManager,
+        moduleDescriptor: ModuleDescriptor,
+        configuration: DeserializationConfiguration
+
+    ):  PackageFragmentProvider {
+        val deserializedPackageFragments = packageFragmentsFactory.createCachedPackageFragments(
+            byteArrays, moduleDescriptor, storageManager
+        )
+
+        val provider = PackageFragmentProviderImpl(deserializedPackageFragments)
+        return initializePackageFragmentProvider(provider, deserializedPackageFragments, storageManager,
+            moduleDescriptor, configuration, null)
+    }
+
     override fun createPackageFragmentProvider(
         library: KotlinLibrary,
         packageAccessHandler: PackageAccessHandler?,
@@ -98,6 +114,18 @@ class KlibMetadataModuleDescriptorFactoryImpl(
         )
 
         val provider = PackageFragmentProviderImpl(deserializedPackageFragments + syntheticPackageFragments)
+        return initializePackageFragmentProvider(provider, deserializedPackageFragments, storageManager,
+            moduleDescriptor, configuration, compositePackageFragmentAddend)
+    }
+
+    fun initializePackageFragmentProvider(
+        provider: PackageFragmentProviderImpl,
+        fragmentsToInitialize: List<DeserializedPackageFragment>,
+        storageManager: StorageManager,
+        moduleDescriptor: ModuleDescriptor,
+        configuration: DeserializationConfiguration,
+        compositePackageFragmentAddend: PackageFragmentProvider?
+    ): PackageFragmentProvider {
 
         val notFoundClasses = NotFoundClasses(storageManager, moduleDescriptor)
 
@@ -124,8 +152,8 @@ class KlibMetadataModuleDescriptorFactoryImpl(
             extensionRegistryLite = KlibMetadataSerializerProtocol.extensionRegistry
         )
 
-        for (packageFragment in deserializedPackageFragments) {
+        fragmentsToInitialize.forEach {
-            packageFragment.initialize(components)
+            it.initialize(components)
         }
 
         return compositePackageFragmentAddend?.let {