[klibs] Added deserialization strategy: on-demand

With on-demand strategy only file header is deserialized
2022-04-25 11:08:26 +05:00
parent 025a21761b
commit 31eedebec6
4 changed files with 66 additions and 118 deletions
@@ -28,8 +28,7 @@ abstract class BasicIrModuleDeserializer(
    override val strategyResolver: (String) -> DeserializationStrategy,
    libraryAbiVersion: KotlinAbiVersion,
    private val containsErrorCode: Boolean = false
-) :
-    IrModuleDeserializer(moduleDescriptor, libraryAbiVersion) {
+) : IrModuleDeserializer(moduleDescriptor, libraryAbiVersion) {

    private val fileToDeserializerMap = mutableMapOf<IrFile, IrFileDeserializer>()

@@ -38,25 +37,34 @@ abstract class BasicIrModuleDeserializer(
    val moduleReversedFileIndex = mutableMapOf<IdSignature, FileDeserializationState>()

    override val moduleDependencies by lazy {
-        moduleDescriptor.allDependencyModules.filter { it != moduleDescriptor }.map { linker.resolveModuleDeserializer(it, null) }
+        moduleDescriptor.allDependencyModules
+            .filter { it != moduleDescriptor }
+            .map { linker.resolveModuleDeserializer(it, null) }
    }

    override fun fileDeserializers(): Collection<IrFileDeserializer> {
-        return fileToDeserializerMap.values
+        return fileToDeserializerMap.values.filterNot { strategyResolver(it.file.fileEntry.name).onDemand }
    }

+    protected lateinit var fileDeserializationStates: List<FileDeserializationState>
+
    override fun init(delegate: IrModuleDeserializer) {
        val fileCount = klib.fileCount()

-        val files = ArrayList<IrFile>(fileCount)
+        val fileDeserializationStates = mutableListOf<FileDeserializationState>()

        for (i in 0 until fileCount) {
            val fileStream = klib.file(i).codedInputStream
            val fileProto = ProtoFile.parseFrom(fileStream, ExtensionRegistryLite.newInstance())
-            files.add(deserializeIrFile(fileProto, i, delegate, containsErrorCode))
+            val fileReader = IrLibraryFileFromBytes(IrKlibBytesSource(klib, i))
+            val file = fileReader.createFile(moduleFragment, fileProto)
+
+            fileDeserializationStates.add(deserializeIrFile(fileProto, file, fileReader, i, delegate, containsErrorCode))
+            if (!strategyResolver(file.fileEntry.name).onDemand)
+                moduleFragment.files.add(file)
        }

-        moduleFragment.files.addAll(files)
+        this.fileDeserializationStates = fileDeserializationStates

        fileToDeserializerMap.values.forEach { it.symbolDeserializer.deserializeExpectActualMapping() }
    }
@@ -89,10 +97,9 @@ abstract class BasicIrModuleDeserializer(
    // TODO: fix to topLevel checker
    override fun contains(idSig: IdSignature): Boolean = idSig in moduleReversedFileIndex

-    override fun deserializeIrSymbol(idSig: IdSignature, symbolKind: BinarySymbolData.SymbolKind): IrSymbol {
+    open fun tryDeserializeIrSymbol(idSig: IdSignature, symbolKind: BinarySymbolData.SymbolKind): IrSymbol? {
        val topLevelSignature = idSig.topLevelSignature()
-        val fileLocalDeserializationState = moduleReversedFileIndex[topLevelSignature]
-            ?: error("No file for $topLevelSignature (@ $idSig) in module $moduleDescriptor")
+        val fileLocalDeserializationState = moduleReversedFileIndex[topLevelSignature] ?: return null

        fileLocalDeserializationState.addIdSignature(topLevelSignature)
        moduleDeserializationState.enqueueFile(fileLocalDeserializationState)
@@ -100,12 +107,16 @@ abstract class BasicIrModuleDeserializer(
        return fileLocalDeserializationState.fileDeserializer.symbolDeserializer.deserializeIrSymbol(idSig, symbolKind)
    }

+    final override fun deserializeIrSymbol(idSig: IdSignature, symbolKind: BinarySymbolData.SymbolKind) =
+        tryDeserializeIrSymbol(idSig, symbolKind)
+            ?: error("No file for ${idSig.topLevelSignature()} (@ $idSig) in module $moduleDescriptor")
+
    override val moduleFragment: IrModuleFragment = IrModuleFragmentImpl(moduleDescriptor, linker.builtIns, emptyList())

-    private fun deserializeIrFile(fileProto: ProtoFile, fileIndex: Int, moduleDeserializer: IrModuleDeserializer, allowErrorNodes: Boolean): IrFile {
-
-        val fileReader = IrLibraryFileFromBytes(IrKlibBytesSource(moduleDeserializer.klib, fileIndex))
-        val file = fileReader.createFile(moduleFragment, fileProto)
+    private fun deserializeIrFile(
+        fileProto: ProtoFile, file: IrFile, fileReader: IrLibraryFileFromBytes,
+        fileIndex: Int, moduleDeserializer: IrModuleDeserializer, allowErrorNodes: Boolean
+    ): FileDeserializationState {
        val fileStrategy = strategyResolver(file.fileEntry.name)

        val fileDeserializationState = FileDeserializationState(
@@ -122,19 +133,21 @@ abstract class BasicIrModuleDeserializer(

        fileToDeserializerMap[file] = fileDeserializationState.fileDeserializer

-        val topLevelDeclarations = fileDeserializationState.fileDeserializer.reversedSignatureIndex.keys
-        topLevelDeclarations.forEach {
-            moduleReversedFileIndex.putIfAbsent(it, fileDeserializationState) // TODO Why not simple put?
+        if (!fileStrategy.onDemand) {
+            val topLevelDeclarations = fileDeserializationState.fileDeserializer.reversedSignatureIndex.keys
+            topLevelDeclarations.forEach {
+                moduleReversedFileIndex.putIfAbsent(it, fileDeserializationState) // TODO Why not simple put?
+            }
+
+            if (fileStrategy.theWholeWorld) {
+                fileDeserializationState.enqueueAllDeclarations()
+            }
+            if (fileStrategy.theWholeWorld || fileStrategy.explicitlyExported) {
+                moduleDeserializationState.enqueueFile(fileDeserializationState)
+            }
        }

-        if (fileStrategy.theWholeWorld) {
-            fileDeserializationState.enqueueAllDeclarations()
-        }
-        if (fileStrategy.theWholeWorld || fileStrategy.explicitlyExported) {
-            moduleDeserializationState.enqueueFile(fileDeserializationState)
-        }
-
-        return file
+        return fileDeserializationState
    }

    override fun addModuleReachableTopLevel(idSig: IdSignature) {
@@ -58,7 +58,7 @@ class IrFileDeserializer(
 class FileDeserializationState(
    val linker: KotlinIrLinker,
    val fileIndex: Int,
-    file: IrFile,
+    val file: IrFile,
    val fileReader: IrLibraryFileFromBytes,
    fileProto: ProtoFile,
    deserializeBodies: Boolean,
@@ -88,7 +88,7 @@ class FileDeserializationState(
            actualModuleDeserializer.deserializeIrSymbol(idSig, symbolKind)
        }

-    private val declarationDeserializer = IrDeclarationDeserializer(
+    val declarationDeserializer = IrDeclarationDeserializer(
        linker.builtIns,
        linker.symbolTable,
        linker.symbolTable.irFactory,
@@ -467,14 +467,16 @@ abstract class KotlinIrLinker(
 }

 enum class DeserializationStrategy(
+    val onDemand: Boolean,
    val needBodies: Boolean,
    val explicitlyExported: Boolean,
    val theWholeWorld: Boolean,
    val inlineBodies: Boolean
 ) {
-    ONLY_REFERENCED(true, false, false, true),
-    ALL(true, true, true, true),
-    EXPLICITLY_EXPORTED(true, true, false, true),
-    ONLY_DECLARATION_HEADERS(false, false, false, false),
-    WITH_INLINE_BODIES(false, false, false, true)
+    ON_DEMAND(true, true, false, false, true),
+    ONLY_REFERENCED(false, true, false, false, true),
+    ALL(false, true, true, true, true),
+    EXPLICITLY_EXPORTED(false, true, true, false, true),
+    ONLY_DECLARATION_HEADERS(false, false, false, false, false),
+    WITH_INLINE_BODIES(false, false, false, false, true)
 }
@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.backend.konan.serialization
 import org.jetbrains.kotlin.backend.common.lower.parents
 import org.jetbrains.kotlin.backend.common.overrides.FakeOverrideBuilder
 import org.jetbrains.kotlin.backend.common.overrides.FakeOverrideClassFilter
-import org.jetbrains.kotlin.backend.common.overrides.FileLocalAwareLinker
 import org.jetbrains.kotlin.backend.common.serialization.*
 import org.jetbrains.kotlin.backend.common.serialization.encodings.BinaryNameAndType
 import org.jetbrains.kotlin.backend.common.serialization.encodings.BinarySymbolData
@@ -39,7 +38,6 @@ import org.jetbrains.kotlin.descriptors.konan.klibModuleOrigin
 import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.builders.TranslationPluginContext
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.impl.IrFactoryImpl
 import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
 import org.jetbrains.kotlin.ir.declarations.lazy.IrLazyClass
 import org.jetbrains.kotlin.ir.expressions.IrBody
@@ -55,13 +53,11 @@ import org.jetbrains.kotlin.library.KotlinLibrary
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import sun.misc.Unsafe
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrClass as ProtoClass
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrDeclaration as ProtoDeclaration
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrField as ProtoField
-import org.jetbrains.kotlin.backend.common.serialization.proto.IrFile as ProtoFile
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrFunction as ProtoFunction
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrProperty as ProtoProperty

@@ -400,7 +396,7 @@ internal class KonanIrLinker(
            moduleDescriptor: ModuleDescriptor,
            klib: KotlinLibrary,
            strategyResolver: (String) -> DeserializationStrategy
-    ): BasicIrModuleDeserializer(this@KonanIrLinker, moduleDescriptor, klib, strategyResolver, klib.versions.abiVersion ?: KotlinAbiVersion.CURRENT) {
+    ) : BasicIrModuleDeserializer(this@KonanIrLinker, moduleDescriptor, klib, strategyResolver, klib.versions.abiVersion ?: KotlinAbiVersion.CURRENT) {
        override val moduleFragment: IrModuleFragment = KonanIrModuleFragmentImpl(moduleDescriptor, builtIns, emptyList())

        fun buildInlineFunctionReference(irFunction: IrFunction): SerializedInlineFunctionReference {
@@ -620,7 +616,11 @@ internal class KonanIrLinker(
    inner class KonanCachedLibraryModuleDeserializer(
            moduleDescriptor: ModuleDescriptor,
            override val klib: KotlinLibrary
-    ) : IrModuleDeserializer(moduleDescriptor, klib.versions.abiVersion ?: KotlinAbiVersion.CURRENT) {
+    ) : BasicIrModuleDeserializer(this@KonanIrLinker, moduleDescriptor, klib,
+            { _ -> DeserializationStrategy.ON_DEMAND },
+            klib.versions.abiVersion ?: KotlinAbiVersion.CURRENT
+    ) {
+        override val moduleFragment: IrModuleFragment = KonanIrModuleFragmentImpl(moduleDescriptor, builtIns, emptyList())

        private val descriptorByIdSignatureFinder = DescriptorByIdSignatureFinderImpl(
                moduleDescriptor, KonanManglerDesc,
@@ -630,81 +630,16 @@ internal class KonanIrLinker(
        override fun contains(idSig: IdSignature) =
                idSig.isPubliclyVisible && descriptorByIdSignatureFinder.findDescriptorBySignature(idSig) != null

-        override fun deserializeIrSymbol(idSig: IdSignature, symbolKind: BinarySymbolData.SymbolKind): IrSymbol {
+        override fun tryDeserializeIrSymbol(idSig: IdSignature, symbolKind: BinarySymbolData.SymbolKind): IrSymbol {
            val descriptor = descriptorByIdSignatureFinder.findDescriptorBySignature(idSig)
                    ?: error("Expecting descriptor for $idSig")

            return (stubGenerator.generateMemberStub(descriptor) as IrSymbolOwner).symbol
        }

-        override val moduleFragment: IrModuleFragment = KonanIrModuleFragmentImpl(moduleDescriptor, builtIns)
-
-        override val moduleDependencies by lazy {
-            moduleDescriptor.allDependencyModules
-                    .filter { it != moduleDescriptor }
-                    .map { resolveModuleDeserializer(it, null) }
-        }
-
-        inner class FileDeserializationInfo(val fileReader: IrLibraryFileFromBytes, val file: IrFile,
-                                            val declarationDeserializer: IrDeclarationDeserializer,
-                                            val fakeOverrideBuilder: FakeOverrideBuilder)
-
-        private val filesDeserializationInfo by lazy {
-            val result = mutableListOf<FileDeserializationInfo>()
-            val fileCount = klib.fileCount()
-
-            for (i in 0 until fileCount) {
-                val fileStream = klib.file(i).codedInputStream
-                val fileProto = ProtoFile.parseFrom(fileStream, ExtensionRegistryLite.newInstance())
-
-                val fileReader = IrLibraryFileFromBytes(IrKlibBytesSource(klib, i))
-                val file = fileReader.createFile(moduleFragment, fileProto)
-
-                val symbolDeserializer = IrSymbolDeserializer(
-                        symbolTable, fileReader, file.symbol, emptyList(), { }, { _, symbol -> symbol })
-                { idSig, symbolKind ->
-                    val topLevelSig = idSig.topLevelSignature()
-                    val actualModuleDeserializer = resolveModuleDeserializer(moduleDescriptor, null)
-                            .findModuleDeserializerForTopLevelId(topLevelSig)
-                            ?: handleSignatureIdNotFoundInModuleWithDependencies(idSig, this)
-
-                    actualModuleDeserializer.deserializeIrSymbol(idSig, symbolKind)
-                }
-
-                val fakeOverrideBuilder = FakeOverrideBuilder(
-                        object : FileLocalAwareLinker {
-                            override fun tryReferencingSimpleFunctionByLocalSignature(parent: IrDeclaration, idSignature: IdSignature) =
-                                    if (idSignature.isPubliclyVisible) null else symbolDeserializer.referenceSimpleFunctionByLocalSignature(idSignature)
-
-                            override fun tryReferencingPropertyByLocalSignature(parent: IrDeclaration, idSignature: IdSignature) =
-                                    if (idSignature.isPubliclyVisible) null else symbolDeserializer.referencePropertyByLocalSignature(idSignature)
-                        },
-                        symbolTable, KonanManglerIr, IrTypeSystemContextImpl(builtIns), emptyMap(), KonanFakeOverrideClassFilter)
-
-                val declarationDeserializer = IrDeclarationDeserializer(
-                        builtIns,
-                        symbolTable,
-                        IrFactoryImpl,
-                        fileReader,
-                        file,
-                        allowErrorNodes = false,
-                        deserializeInlineFunctions = true,
-                        deserializeBodies = true,
-                        symbolDeserializer,
-                        fakeOverrideBuilder.platformSpecificClassFilter,
-                        fakeOverrideBuilder,
-                        compatibilityMode = CompatibilityMode(libraryAbiVersion)
-                )
-
-                result.add(FileDeserializationInfo(fileReader, file, declarationDeserializer, fakeOverrideBuilder))
-            }
-
-            result
-        }
-
        private val inlineFunctionReferences by lazy {
            cachedLibraries.getLibraryCache(klib)!!.serializedInlineFunctionBodies.associateBy {
-                filesDeserializationInfo[it.file].declarationDeserializer.symbolDeserializer.deserializeIdSignature(it.functionSignature)
+                fileDeserializationStates[it.file].declarationDeserializer.symbolDeserializer.deserializeIdSignature(it.functionSignature)
            }
        }

@@ -727,8 +662,8 @@ internal class KonanIrLinker(
                    ?: error("No signature for ${function.render()}")
            val inlineFunctionReference = inlineFunctionReferences[signature]
                    ?: error("No inline function reference for ${function.render()}, sig = ${signature.render()}")
-            val fileDeserializationInfo = filesDeserializationInfo[inlineFunctionReference.file]
-            val declarationDeserializer = fileDeserializationInfo.declarationDeserializer
+            val fileDeserializationState = fileDeserializationStates[inlineFunctionReference.file]
+            val declarationDeserializer = fileDeserializationState.declarationDeserializer
            val symbolDeserializer = declarationDeserializer.symbolDeserializer

            val outerClasses = (function.parent as? IrClass)?.getOuterClasses(takeOnlyInner = true) ?: emptyList()
@@ -774,21 +709,21 @@ internal class KonanIrLinker(
                    parameter.defaultValue = declarationDeserializer.deserializeExpressionBody(defaultValueIndex)
            }

-            fileDeserializationInfo.fakeOverrideBuilder.provideFakeOverrides()
+            fakeOverrideBuilder.provideFakeOverrides()

            inlineFunctionFiles[packageFragment]?.let {
-                require(it == fileDeserializationInfo.file) {
-                    "Different files ${it.fileEntry.name} and ${fileDeserializationInfo.file.fileEntry.name} have the same $packageFragment"
+                require(it == fileDeserializationState.file) {
+                    "Different files ${it.fileEntry.name} and ${fileDeserializationState.file.fileEntry.name} have the same $packageFragment"
                }
            }
-            inlineFunctionFiles[packageFragment] = fileDeserializationInfo.file
+            inlineFunctionFiles[packageFragment] = fileDeserializationState.file

-            return InlineFunctionOriginInfo(function, fileDeserializationInfo.file, inlineFunctionReference.startOffset, inlineFunctionReference.endOffset)
+            return InlineFunctionOriginInfo(function, fileDeserializationState.file, inlineFunctionReference.startOffset, inlineFunctionReference.endOffset)
        }

        private val classesFields by lazy {
            cachedLibraries.getLibraryCache(klib)!!.serializedClassFields.associateBy {
-                filesDeserializationInfo[it.file].declarationDeserializer.symbolDeserializer.deserializeIdSignature(it.classSignature)
+                fileDeserializationStates[it.file].declarationDeserializer.symbolDeserializer.deserializeIdSignature(it.classSignature)
            }
        }

@@ -797,8 +732,8 @@ internal class KonanIrLinker(
                    ?: error("No signature for ${irClass.render()}")
            val serializedClassFields = classesFields[signature]
                    ?: error("No class fields for ${irClass.render()}, sig = ${signature.render()}")
-            val fileDeserializationInfo = filesDeserializationInfo[serializedClassFields.file]
-            val declarationDeserializer = fileDeserializationInfo.declarationDeserializer
+            val fileDeserializationState = fileDeserializationStates[serializedClassFields.file]
+            val declarationDeserializer = fileDeserializationState.declarationDeserializer
            val symbolDeserializer = declarationDeserializer.symbolDeserializer

            val outerClasses = irClass.getOuterClasses(takeOnlyInner = true)
@@ -825,7 +760,7 @@ internal class KonanIrLinker(
                    require(outerThisField != null) { "For an inner class ${irClass.render()} there should be <outer this> field" }
                    outerThisField.toFieldInfo()
                } else {
-                    val name = fileDeserializationInfo.fileReader.string(field.name)
+                    val name = fileDeserializationState.fileReader.string(field.name)
                    val type = when {
                        field.type != InvalidIndex -> declarationDeserializer.deserializeIrType(field.type)
                        field.binaryType == InvalidIndex -> builtIns.anyNType
@@ -847,8 +782,6 @@ internal class KonanIrLinker(
                }
            }
        }
-
-        override val kind get() = IrModuleDeserializerKind.DESERIALIZED
    }

    private inner class KonanForwardDeclarationModuleDeserializer(moduleDescriptor: ModuleDescriptor) : IrModuleDeserializer(moduleDescriptor, KotlinAbiVersion.CURRENT) {