Load metadata lazily per package fragment.

2017-03-29 19:41:27 +03:00
parent c9c9aee7bb
commit 742934a1d8
5 changed files with 115 additions and 54 deletions
@@ -58,6 +58,11 @@ import java.util.*
 import kotlin.LazyThreadSafetyMode.PUBLICATION
 import kotlin.reflect.KProperty

+internal class LinkData(
+    val module: String,
+    val fragments: List<String>,
+    val fragmentNames: List<String> )
+
 internal class SpecialDescriptorsFactory(val context: Context) {
    private val enumSpecialDescriptorsFactory by lazy { EnumSpecialDescriptorsFactory(context) }
    private val outerThisDescriptors = mutableMapOf<ClassDescriptor, PropertyDescriptor>()
@@ -230,7 +235,7 @@ internal class Context(config: KonanConfig) : KonanBackendContext(config) {
    // right after the frontend.
    // But we have to wait until the code generation phase,
    // to dump this information into generated file.
-    var serializedLinkData: String? = null
+    var serializedLinkData: LinkData? = null

    // TODO: make lateinit?
    var irModule: IrModuleFragment? = null
@@ -35,35 +35,46 @@ fun loadMetadata(configuration: CompilerConfiguration, file: File): ModuleDescri
    var moduleName: String? = null
    val currentAbiVersion = configuration.get(KonanConfigKeys.ABI_VERSION)

-    reader.use {
-        val (nodeCount, kmetadataNodeArg) = reader.namedMetadataNode("kmetadata")
+    val (nodeCount, kmetadataNodeArg) = reader.namedMetadataNode("kmetadata")

-        if (nodeCount != 1) {
-            throw Error("Unknown metadata format. The 'kmetadata' node has ${nodeCount} arguments. Don't know what to do.")
-        }
-
-        val operands = reader.metadataNodeOperands(kmetadataNodeArg)
-
-        val abiNode = operands[0]
-        val nameNode = operands[1]
-        val dataNode = operands[2]
-
-        val abiVersion = reader.string(abiNode).toInt()
-
-        if (abiVersion != currentAbiVersion) {
-            throw Error("Expected ABI version ${currentAbiVersion}, but the binary is ${abiVersion}")
-        }
-        moduleName = reader.string(nameNode)
-
-        metadataAsString = reader.string(dataNode)
+    if (nodeCount != 1) {
+        throw Error("Unknown metadata format. The 'kmetadata' node has ${nodeCount} arguments. Don't know what to do.")
    }

+    val operands = reader.metadataNodeOperands(kmetadataNodeArg)
+
+    val abiNode = operands[0]
+    val nameNode = operands[1]
+    val dataNode = operands[2]
+
+    val abiVersion = reader.string(abiNode).toInt()
+
+    if (abiVersion != currentAbiVersion) {
+        throw Error("Expected ABI version ${currentAbiVersion}, but the binary is ${abiVersion}")
+    }
+    moduleName = reader.string(nameNode)
+
+    val tableOfContentsAsString = reader.string(dataNode)
+
    val moduleDescriptor = 
-        deserializeModule(configuration, metadataAsString!!, moduleName!!)
+        deserializeModule(configuration, {it->loadPackageFragment(reader, it)}, tableOfContentsAsString!!, moduleName!!)

    return moduleDescriptor
 }

+fun loadPackageFragment(reader: MetadataReader, fqName: String): String {
+    val (nodeCount, kpackageNodeArg) = reader.namedMetadataNode("kpackage:$fqName")
+
+    if (nodeCount != 1) {
+        throw Error("The 'kpackage:$fqName' node has ${nodeCount} arguments.")
+    }
+
+    val operands = reader.metadataNodeOperands(kpackageNodeArg)
+    val dataNode = operands[0]
+    val base64 =  reader.string(dataNode)
+    return base64
+}
+
 class MetadataReader(file: File) : Closeable {

    lateinit var llvmModule: LLVMModuleRef
@@ -165,16 +176,26 @@ internal class MetadataGenerator(override val context: Context): ContextUtils {
    }

    private fun addLinkData(module: IrModuleFragment) {
-        val moduleAsString = context.serializedLinkData
-        if (moduleAsString == null) return
+        val linkData = context.serializedLinkData
+        if (linkData == null) return

        val abiVersion = context.config.configuration.get(KonanConfigKeys.ABI_VERSION)
        val abiNode = metadataString("$abiVersion")
        val moduleName = metadataString(module.descriptor.name.asString())
-        val dataNode = metadataString(moduleAsString)
+        val module = linkData.module
+        val fragments = linkData.fragments
+        val fragmentNames = linkData.fragmentNames
+        val dataNode = metadataString(module)

        val kmetadataArg  = metadataNode(listOf(abiNode, moduleName, dataNode))
        emitModuleMetadata("kmetadata", kmetadataArg)
+
+        fragments.forEachIndexed { index, it ->
+            val name = fragmentNames.get(index)
+            val dataNode = metadataString(it)
+            val kpackageArg = metadataNode(listOf(dataNode))
+            emitModuleMetadata("kpackage:$name", kpackageArg)
+        }
    }

    internal fun endModule(module: IrModuleFragment) {
@@ -107,6 +107,6 @@ message Classes {
 }

 message Library {
-    repeated PackageFragment package_fragment = 2;
+    repeated string package_fragment_name = 1;
 }

@@ -25,15 +25,25 @@ import org.jetbrains.kotlin.serialization.deserialization.NameResolverImpl
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPackageMemberScope
 import org.jetbrains.kotlin.storage.StorageManager

-class KonanPackageFragment( val proto: KonanLinkData.PackageFragment,
+class KonanPackageFragment(fqNameString: String,
+    packageLoader: (String)->KonanLinkData.PackageFragment,
    storageManager: StorageManager, module: ModuleDescriptor) : 
-    DeserializedPackageFragment(FqName(proto.getFqName()), 
+    DeserializedPackageFragment(FqName(fqNameString), 
        storageManager, module) {

-    private val nameResolver = 
-        NameResolverImpl(proto.getStringTable(), proto.getNameTable())
+    // The proto field is lazy so that we can load only needed
+    // packages from the library.
+    val proto: KonanLinkData.PackageFragment by lazy {
+        packageLoader(fqNameString)
+    }

-    override val classDataFinder = KonanClassDataFinder(proto, nameResolver)
+    private val nameResolver by lazy {
+        NameResolverImpl(proto.getStringTable(), proto.getNameTable())
+    }
+
+    override val classDataFinder by lazy {
+        KonanClassDataFinder(proto, nameResolver)
+    }

    override fun computeMemberScope(): DeserializedPackageMemberScope {
        val packageProto = proto.getPackage()
@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.backend.konan.serialization

 import org.jetbrains.kotlin.backend.konan.Context
+import org.jetbrains.kotlin.backend.konan.LinkData
 import org.jetbrains.kotlin.backend.konan.KonanBuiltIns
 import org.jetbrains.kotlin.backend.konan.llvm.base64Decode
 import org.jetbrains.kotlin.backend.konan.llvm.base64Encode
@@ -48,9 +49,28 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.SimpleType
 import java.util.zip.GZIPInputStream
 import java.util.zip.GZIPOutputStream
+import java.io.InputStream
 import java.io.ByteArrayInputStream
 import java.io.ByteArrayOutputStream

+typealias Base64 = String
+
+fun byteArrayToBase64(byteArray: ByteArray): Base64 {
+    val gzipped = ByteArrayOutputStream()
+    val gzipStream = GZIPOutputStream(gzipped)
+    gzipStream.write(byteArray)
+    gzipStream.close()
+    val base64 = base64Encode(gzipped.toByteArray())
+    return base64
+}
+
+fun base64ToStream(base64: Base64): InputStream {
+    val gzipped = base64Decode(base64)
+    return GZIPInputStream(ByteArrayInputStream(gzipped))
+
+}
+
+
 /* ------------ Deserializer part ------------------------------------------*/

 object NullFlexibleTypeDeserializer : FlexibleTypeDeserializer {
@@ -61,12 +81,13 @@ object NullFlexibleTypeDeserializer : FlexibleTypeDeserializer {
 }

 fun createKonanPackageFragmentProvider(
-        fragments: List<KonanLinkData.PackageFragment>, 
+        fragmentNames: List<String>, 
+        packageLoader: (String)->KonanLinkData.PackageFragment,
        storageManager: StorageManager, module: ModuleDescriptor, 
        configuration: DeserializationConfiguration): PackageFragmentProvider {

-    val packageFragments = fragments.map { 
-        KonanPackageFragment(it, storageManager, module) 
+    val packageFragments = fragmentNames.map{ 
+        KonanPackageFragment(it, packageLoader, storageManager, module) 
    }
    val provider = PackageFragmentProviderImpl(packageFragments)

@@ -91,8 +112,14 @@ fun createKonanPackageFragmentProvider(
    return provider
 }

+internal fun deserializePackageFragment(base64: Base64): KonanLinkData.PackageFragment {
+    return KonanLinkData.PackageFragment
+        .parseFrom(base64ToStream(base64), 
+            KonanSerializerProtocol.extensionRegistry)
+}
+
 internal fun deserializeModule(configuration: CompilerConfiguration, 
-    base64: String, moduleName: String): ModuleDescriptorImpl {
+    packageLoader:(String)->Base64, library: Base64,  moduleName: String): ModuleDescriptorImpl {

    val storageManager = LockBasedStorageManager()
    val builtIns = KonanBuiltIns(storageManager)
@@ -101,13 +128,14 @@ internal fun deserializeModule(configuration: CompilerConfiguration,
    builtIns.builtInsModule = moduleDescriptor
    val deserializationConfiguration = CompilerDeserializationConfiguration(configuration)

-    val gzipped = base64Decode(base64)
-    val libraryAsByteArray = GZIPInputStream(ByteArrayInputStream(gzipped))
-    val libraryProto = KonanLinkData.Library.parseFrom(
-        libraryAsByteArray, KonanSerializerProtocol.extensionRegistry)
+    val libraryProto = KonanLinkData.Library
+        .parseFrom(base64ToStream(library), 
+            KonanSerializerProtocol.extensionRegistry)

    val provider = createKonanPackageFragmentProvider(
-        libraryProto.packageFragmentList, storageManager, 
+        libraryProto.packageFragmentNameList,
+        {it -> deserializePackageFragment(packageLoader(it))},
+        storageManager, 
        moduleDescriptor, deserializationConfiguration)

    moduleDescriptor.initialize(provider)
@@ -215,24 +243,21 @@ internal class KonanSerializationUtil(val context: Context) {
            }
        }
    }
-
-    internal fun serializeModule(moduleDescriptor: ModuleDescriptor): String {
-        val library = KonanLinkData.Library.newBuilder()
-
+    internal fun serializeModule(moduleDescriptor: ModuleDescriptor): LinkData {
+        val libraryProto = KonanLinkData.Library.newBuilder()
+        val fragments = mutableListOf<String>()
+        val fragmentNames = mutableListOf<String>()

        getPackagesFqNames(moduleDescriptor).forEach {
            val packageProto = serializePackage(it, moduleDescriptor)
-            library.addPackageFragment(packageProto)
+            libraryProto.addPackageFragmentName(it.asString())
+            fragments.add(
+                byteArrayToBase64(packageProto.toByteArray()))
+            fragmentNames.add(it.asString())
        }
-
-        val libraryAsByteArray = library.build().toByteArray()
-
-        val gzipped = ByteArrayOutputStream()
-        val gzipStream = GZIPOutputStream(gzipped)
-        gzipStream.write(libraryAsByteArray)
-        gzipStream.close()
-        val base64 = base64Encode(gzipped.toByteArray())
-        return base64
+        val libraryAsByteArray = libraryProto.build().toByteArray()
+        val library = byteArrayToBase64(libraryAsByteArray)
+        return LinkData(library, fragments, fragmentNames)
    }

    /* The section below is specific for IR serialization */