[JS KLIB IC] Support incremental klib generation for JS IR

Note: the following test do not work yet because of IrField FO issue
testSecondaryConstructorAdded
testCompanionObjectInheritedMemberChanged

compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/klib.kt

@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.ir.backend.js
 
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VfsUtilCore
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.backend.common.LoggingContext
 import org.jetbrains.kotlin.backend.common.serialization.DescriptorTable
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
@@ -14,13 +16,14 @@ import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.config.languageVersionSettings
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.impl.CompositePackageFragmentProvider
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.incremental.components.LookupTracker
-import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.*
 import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsIrLinker
 import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsIrModuleSerializer
+import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsMangler
 import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.newJsDescriptorUniqId
 import org.jetbrains.kotlin.ir.backend.js.lower.serialization.metadata.*
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
@@ -30,12 +33,14 @@ import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
 import org.jetbrains.kotlin.js.analyze.TopDownAnalyzerFacadeForJS
 import org.jetbrains.kotlin.js.analyzer.JsAnalysisResult
+import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 import org.jetbrains.kotlin.konan.KonanVersionImpl
 import org.jetbrains.kotlin.konan.MetaVersion
 import org.jetbrains.kotlin.konan.properties.propertyList
 import org.jetbrains.kotlin.library.*
 import org.jetbrains.kotlin.library.impl.buildKoltinLibrary
 import org.jetbrains.kotlin.library.impl.createKotlinLibrary
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 import org.jetbrains.kotlin.psi.KtFile
@@ -43,9 +48,11 @@ import org.jetbrains.kotlin.psi2ir.Psi2IrConfiguration
 import org.jetbrains.kotlin.psi2ir.Psi2IrTranslator
 import org.jetbrains.kotlin.psi2ir.generators.GeneratorContext
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.BindingContextUtils
 import org.jetbrains.kotlin.resolve.CompilerDeserializationConfiguration
 import org.jetbrains.kotlin.storage.LockBasedStorageManager
 import org.jetbrains.kotlin.utils.DFS
+import java.io.File
 import org.jetbrains.kotlin.konan.file.File as KFile
 
 
@@ -73,6 +80,8 @@ private val emptyLoggingContext = object : LoggingContext {
 private val CompilerConfiguration.metadataVersion
     get() = get(CommonConfigurationKeys.METADATA_VERSION) as? JsKlibMetadataVersion ?: JsKlibMetadataVersion.INSTANCE
 
+class KotlinFileSerializedData(val metadata: ByteArray, val irData: SerializedIrFile)
+
 fun generateKLib(
     project: Project,
     files: List<KtFile>,
@@ -82,12 +91,39 @@ fun generateKLib(
     outputKlibPath: String,
     nopack: Boolean
 ) {
+    val incrementalDataProvider = configuration.get(JSConfigurationKeys.INCREMENTAL_DATA_PROVIDER)
+
+    val icData: List<KotlinFileSerializedData>
+
+    if (incrementalDataProvider != null) {
+        val nonCompiledSources = files.map { VfsUtilCore.virtualToIoFile(it.virtualFile) to it }.toMap()
+        val compiledIrFiles = incrementalDataProvider.serializedIrFiles
+        val compiledMetaFiles = incrementalDataProvider.compiledPackageParts
+
+        assert(compiledIrFiles.size == compiledMetaFiles.size)
+
+        val storage = mutableListOf<KotlinFileSerializedData>()
+
+        for (f in compiledIrFiles.keys) {
+            if (f in nonCompiledSources) continue
+
+            val irData = compiledIrFiles[f] ?: error("No Ir Data found for file $f")
+            val metaFile = compiledMetaFiles[f] ?: error("No Meta Data found for file $f")
+            val irFile = with(irData) {
+                SerializedIrFile(fileData, String(fqn), f.path, symbols, types, strings, bodies, declarations)
+            }
+            storage.add(KotlinFileSerializedData(metaFile.metadata, irFile))
+        }
+
+        icData = storage
+    } else {
+        icData = emptyList()
+    }
+
     val depsDescriptors = ModulesStructure(project, files, configuration, allDependencies, friendDependencies)
 
     val psi2IrContext = runAnalysisAndPreparePsi2Ir(depsDescriptors)
 
-    trySaveIncrementalData(psi2IrContext, configuration, files)
-
     val moduleFragment = psi2IrContext.generateModuleFragment(files)
 
     val moduleName = configuration[CommonConfigurationKeys.MODULE_NAME]!!
@@ -96,12 +132,13 @@ fun generateKLib(
 
     serializeModuleIntoKlib(
         moduleName,
-        configuration.metadataVersion,
+        configuration,
-        configuration.languageVersionSettings,
         psi2IrContext.bindingContext,
+        files,
         outputKlibPath,
         allDependencies,
         moduleFragment,
+        icData,
         nopack
     )
 }
@@ -211,7 +248,6 @@ private fun loadKlibMetadata(
     return md
 }
 
-
 private class ModulesStructure(
     private val project: Project,
     private val files: List<KtFile>,
@@ -284,34 +320,91 @@ private class ModulesStructure(
             null // null in case compiling builtInModule itself
 }
 
+private fun getDescriptorForElement(
+    context: BindingContext,
+    element: PsiElement
+): DeclarationDescriptor = BindingContextUtils.getNotNull(context, BindingContext.DECLARATION_TO_DESCRIPTOR, element)
+
 fun serializeModuleIntoKlib(
     moduleName: String,
-    metadataVersion: JsKlibMetadataVersion,
+    configuration: CompilerConfiguration,
-    languageVersionSettings: LanguageVersionSettings,
     bindingContext: BindingContext,
+    files: List<KtFile>,
     klibPath: String,
     dependencies: List<KotlinLibrary>,
     moduleFragment: IrModuleFragment,
+    cleanFiles: List<KotlinFileSerializedData>,
     nopack: Boolean
 ) {
-    val descriptorTable = DescriptorTable()
+    assert(files.size == moduleFragment.files.size)
 
+    val descriptorTable = DescriptorTable()
     val serializedIr =
         JsIrModuleSerializer(emptyLoggingContext, moduleFragment.irBuiltins, descriptorTable).serializedIrModule(moduleFragment)
-    val serializer = JsKlibMetadataSerializationUtil
 
-    val moduleDescription =
+    val moduleDescriptor = moduleFragment.descriptor
-        JsKlibMetadataModuleDescriptor(moduleName, dependencies.map { it.moduleName }, moduleFragment.descriptor)
+
-    val serializedMetadata = serializer.serializeMetadata(
+    val descriptorSelector = { declarationDescriptor: DeclarationDescriptor ->
-        bindingContext,
-        moduleDescription,
-        languageVersionSettings,
-        metadataVersion
-    ) { declarationDescriptor ->
         val index = descriptorTable.get(declarationDescriptor) ?: error("No descriptor ID found for $declarationDescriptor")
         newJsDescriptorUniqId(index)
     }
 
+    val metadataVersion = configuration.metadataVersion
+    val languageVersionSettings = configuration.languageVersionSettings
+
+    fun serializeScope(fqName: FqName, memberScope: Collection<DeclarationDescriptor>): ByteArray {
+        return JsKlibMetadataSerializationUtil.serializeDescriptors(
+            bindingContext,
+            moduleDescriptor,
+            memberScope,
+            fqName,
+            languageVersionSettings,
+            metadataVersion,
+            descriptorSelector
+        ).toByteArray()
+    }
+
+    val incrementalResultsConsumer = configuration.get(JSConfigurationKeys.INCREMENTAL_RESULTS_CONSUMER)
+    val empty = ByteArray(0)
+
+    fun processCompiledFileData(ioFile: File, compiledFile: KotlinFileSerializedData) {
+        incrementalResultsConsumer?.run {
+            processPackagePart(ioFile, compiledFile.metadata, empty, empty)
+            with(compiledFile.irData) {
+                processIrFile(ioFile, fileData, symbols, types, strings, declarations, bodies, fqName.toByteArray())
+            }
+        }
+    }
+
+    val additionalFiles = mutableListOf<KotlinFileSerializedData>()
+
+    for ((ktFile, binaryFile) in files.zip(serializedIr.files)) {
+        val ioFile = VfsUtilCore.virtualToIoFile(ktFile.virtualFile)
+        assert(ioFile.path == binaryFile.path) { "The Kt and Ir files are put in different order" }
+        val memberScope = ktFile.declarations.map { getDescriptorForElement(bindingContext, it) }
+        val packageFragment = serializeScope(ktFile.packageFqName, memberScope)
+        val compiledKotlinFile = KotlinFileSerializedData(packageFragment, binaryFile)
+
+        additionalFiles += compiledKotlinFile
+        processCompiledFileData(ioFile, compiledKotlinFile)
+    }
+
+    incrementalResultsConsumer?.run {
+        processHeader(JsKlibMetadataSerializationUtil.serializeHeader(moduleDescriptor, null, languageVersionSettings).toByteArray())
+    }
+
+    val compiledKotlinFiles = cleanFiles + additionalFiles
+
+    val serializedMetadata =
+        JsKlibMetadataSerializationUtil.serializedMetadata(
+            moduleDescriptor,
+            languageVersionSettings,
+            dependencies.map { it.moduleName },
+            compiledKotlinFiles.groupBy { it.irData.fqName }.map { (fqn, data) -> fqn to data.sortedBy { it.irData.path }.map { it.metadata } }.toMap()
+        )
+
+    val fullSerializedIr = SerializedIrModule(compiledKotlinFiles.map { it.irData })
+
     val abiVersion = KotlinAbiVersion.CURRENT
     val compilerVersion = KonanVersionImpl(MetaVersion.DEV, 1, 3, 0, -1)
     val libraryVersion = "JSIR"
@@ -320,7 +413,7 @@ fun serializeModuleIntoKlib(
 
     buildKoltinLibrary(
         linkDependencies = dependencies,
-        ir = serializedIr,
+        ir = fullSerializedIr,
         metadata = serializedMetadata,
         dataFlowGraph = null,
         manifestProperties = null,

compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/klibIncremental.kt

@@ -1,66 +0,0 @@
-/*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.ir.backend.js
-
-import com.intellij.openapi.vfs.VfsUtilCore
-import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.config.CommonConfigurationKeys
-import org.jetbrains.kotlin.config.CompilerConfiguration
-import org.jetbrains.kotlin.config.languageVersionSettings
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.ModuleDescriptor
-import org.jetbrains.kotlin.js.config.JSConfigurationKeys.INCREMENTAL_RESULTS_CONSUMER
-import org.jetbrains.kotlin.metadata.ProtoBuf
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi2ir.generators.GeneratorContext
-import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.BindingContextUtils
-import org.jetbrains.kotlin.serialization.js.KotlinJavascriptSerializationUtil
-import org.jetbrains.kotlin.utils.JsMetadataVersion
-
-// todo: think if possible to unify with org.jetbrains.kotlin.js.facade.K2JSTranslator.trySaveIncrementalData
-internal fun trySaveIncrementalData(context: GeneratorContext, config: CompilerConfiguration, files: List<KtFile>) {
-    val incrementalResults = config.get(INCREMENTAL_RESULTS_CONSUMER) ?: return
-
-    val bindingContext = context.bindingContext
-    val moduleDescriptor = context.moduleDescriptor
-
-    val emptyByteArray = emptyArray<Byte>().toByteArray()
-
-    for (file in files) {
-        val memberScope = file.declarations.map { getDescriptorForElement(bindingContext, it) }
-        val packagePart = serializeScope(bindingContext, moduleDescriptor, file.packageFqName, memberScope, config)
-        val ioFile = VfsUtilCore.virtualToIoFile(file.virtualFile)
-        incrementalResults.processPackagePart(ioFile, packagePart.toByteArray(), binaryAst = emptyByteArray, inlineData = emptyByteArray)
-    }
-
-    val settings = config.languageVersionSettings
-    incrementalResults.processHeader(KotlinJavascriptSerializationUtil.serializeHeader(moduleDescriptor, null, settings).toByteArray())
-}
-
-// copied from org.jetbrains.kotlin.js.translate.utils.BindingUtils.getDescriptorForElement
-fun getDescriptorForElement(
-    context: BindingContext,
-    element: PsiElement
-): DeclarationDescriptor = BindingContextUtils.getNotNull(context, BindingContext.DECLARATION_TO_DESCRIPTOR, element)
-
-// copied from org.jetbrains.kotlin.js.facade.K2JSTranslator.serializeScope
-private fun serializeScope(
-    bindingContext: BindingContext,
-    moduleDescriptor: ModuleDescriptor,
-    packageName: FqName,
-    scope: Collection<DeclarationDescriptor>,
-    configuration: CompilerConfiguration
-): ProtoBuf.PackageFragment =
-    KotlinJavascriptSerializationUtil.serializeDescriptors(
-        bindingContext,
-        moduleDescriptor,
-        scope,
-        packageName,
-        configuration.languageVersionSettings,
-        configuration.get(CommonConfigurationKeys.METADATA_VERSION) ?: JsMetadataVersion.INSTANCE
-    )

compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/lower/serialization/metadata/JsKlibMetadataSerializationUtil.kt

@@ -5,8 +5,10 @@
 
 package org.jetbrains.kotlin.ir.backend.js.lower.serialization.metadata
 
+import org.jetbrains.kotlin.config.AnalysisFlags
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.library.SerializedMetadata
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
@@ -16,50 +18,102 @@ import org.jetbrains.kotlin.resolve.checkers.ExpectedActualDeclarationChecker
 import org.jetbrains.kotlin.resolve.descriptorUtil.filterOutSourceAnnotations
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
-import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
-import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.serialization.AnnotationSerializer
 import org.jetbrains.kotlin.serialization.DescriptorSerializer
+import org.jetbrains.kotlin.serialization.StringTableImpl
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedClassDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPropertyDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedSimpleFunctionDescriptor
+import java.io.ByteArrayInputStream
+import java.io.ByteArrayOutputStream
+import java.io.OutputStream
 
 object JsKlibMetadataSerializationUtil {
-    fun serializeMetadata(
+
-        bindingContext: BindingContext,
+    fun serializedMetadata(
-        jsDescriptor: JsKlibMetadataModuleDescriptor<ModuleDescriptor>,
+        module: ModuleDescriptor,
-        languageVersionSettings: LanguageVersionSettings,
+        settings: LanguageVersionSettings,
-        metadataVersion: JsKlibMetadataVersion,
+        importedModules: List<String>,
-        declarationTableHandler: (DeclarationDescriptor) -> JsKlibMetadataProtoBuf.DescriptorUniqId?
+        fragments: Map<String, List<ByteArray>>
     ): SerializedMetadata {
-        val libraryProto = JsKlibMetadataProtoBuf.Library.newBuilder()
-        jsDescriptor.imported.forEach { libraryProto.addImportedModule(it) }
-
-        val serializedFragments = HashMap<FqName, ProtoBuf.PackageFragment>()
-        val module = jsDescriptor.data
-        val fragments = mutableListOf<List<ByteArray>>()
         val fragmentNames = mutableListOf<String>()
+        val fragmentParts = mutableListOf<List<ByteArray>>()
 
-        for (fqName in getPackagesFqNames(module).sortedBy { it.asString() }) {
+        for ((fqName, fragment) in fragments.entries.sortedBy { it.key }) {
-            val fragment = serializeDescriptors(
+            fragmentNames += fqName
-                bindingContext, module,
+            fragmentParts += fragment
-                module.getPackage(fqName).memberScope.getContributedDescriptors(),
+        }
-                fqName, languageVersionSettings, metadataVersion, declarationTableHandler
-            )
 
-            if (!fragment.isEmpty()) {
+        val stream = ByteArrayOutputStream()
-                serializedFragments[fqName] = fragment
+
+        serializeHeader(module, null, settings).writeDelimitedTo(stream)
+        asLibrary().writeTo(stream)
+        stream.appendPackageFragments(fragments)
+        importedModules.forEach {
+            stream.writeProto(JsKlibMetadataProtoBuf.Library.IMPORTED_MODULE_FIELD_NUMBER, it.toByteArray())
+        }
+
+        val moduleLibrary = stream.toByteArray()
+
+        return SerializedMetadata(moduleLibrary, fragmentParts, fragmentNames)
+    }
+
+    private fun asLibrary(): JsKlibMetadataProtoBuf.Library {
+        return JsKlibMetadataProtoBuf.Library.newBuilder().build()
+    }
+
+    private fun OutputStream.writeProto(fieldNumber: Int, content: ByteArray) {
+        // Message header
+        write((fieldNumber shl 3) or 2)
+        // Size varint
+        var size = content.size
+        while (size > 0x7F) {
+            write(0x80 or (size and 0x7F))
+            size = size ushr 7
+        }
+        write(size)
+        // Fragment itself
+        write(content)
+    }
+
+    private fun OutputStream.appendPackageFragments(serializedFragments: Map<String, List<ByteArray>>) {
+        for ((_, fragments) in serializedFragments.entries.sortedBy { it.key }) {
+            for (fragment in fragments) {
+                writeProto(JsKlibMetadataProtoBuf.Library.PACKAGE_FRAGMENT_FIELD_NUMBER, fragment)
             }
         }
+    }
 
-        for ((fqName, fragment) in serializedFragments.entries.sortedBy { (fqName, _) -> fqName.asString() }) {
+    fun serializeHeader(
-            libraryProto.addPackageFragment(fragment)
+        module: ModuleDescriptor, packageFqName: FqName?, languageVersionSettings: LanguageVersionSettings
-            fragments.add(listOf(fragment.toByteArray()))
+    ): JsKlibMetadataProtoBuf.Header {
-            fragmentNames.add(fqName.asString())
+        val header = JsKlibMetadataProtoBuf.Header.newBuilder()
+
+        if (packageFqName != null) {
+            header.packageFqName = packageFqName.asString()
         }
 
-        val libraryAsByteArray = libraryProto.build().toByteArray()
+        if (languageVersionSettings.isPreRelease()) {
-        return SerializedMetadata(libraryAsByteArray, fragments, fragmentNames)
+            header.flags = 1
+        }
+
+        val experimentalAnnotationFqNames = languageVersionSettings.getFlag(AnalysisFlags.experimental)
+        if (experimentalAnnotationFqNames.isNotEmpty()) {
+            val stringTable = StringTableImpl()
+            for (fqName in experimentalAnnotationFqNames) {
+                val descriptor = module.resolveClassByFqName(FqName(fqName), NoLookupLocation.FOR_ALREADY_TRACKED) ?: continue
+                header.addAnnotation(ProtoBuf.Annotation.newBuilder().apply {
+                    id = stringTable.getFqNameIndex(descriptor)
+                })
+            }
+            val (strings, qualifiedNames) = stringTable.buildProto()
+            header.strings = strings
+            header.qualifiedNames = qualifiedNames
+        }
+
+        // TODO: write JS code binary version
+
+        return header.build()
     }
 
     fun serializeDescriptors(
@@ -143,33 +197,11 @@ object JsKlibMetadataSerializationUtil {
         return filesProto.build()
     }
 
-    private fun ProtoBuf.PackageFragment.isEmpty(): Boolean =
-        class_Count == 0 && `package`.let { it.functionCount == 0 && it.propertyCount == 0 && it.typeAliasCount == 0 }
-
-    private fun getPackagesFqNames(module: ModuleDescriptor): Set<FqName> {
-        return mutableSetOf<FqName>().apply {
-            getSubPackagesFqNames(module.getPackage(FqName.ROOT), this)
-            add(FqName.ROOT)
-        }
-    }
-
-    private fun getSubPackagesFqNames(packageView: PackageViewDescriptor, result: MutableSet<FqName>) {
-        val fqName = packageView.fqName
-        if (!fqName.isRoot) {
-            result.add(fqName)
-        }
-
-        for (descriptor in packageView.memberScope.getContributedDescriptors(DescriptorKindFilter.PACKAGES, MemberScope.ALL_NAME_FILTER)) {
-            if (descriptor is PackageViewDescriptor) {
-                getSubPackagesFqNames(descriptor, result)
-            }
-        }
-    }
-
     @JvmStatic
     fun readModuleAsProto(metadata: ByteArray): JsKlibMetadataParts {
-        val header = JsKlibMetadataProtoBuf.Header.parseFrom(metadata, JsKlibMetadataSerializerProtocol.extensionRegistry)
+        val inputStream = ByteArrayInputStream(metadata)
-        val content = JsKlibMetadataProtoBuf.Library.parseFrom(metadata, JsKlibMetadataSerializerProtocol.extensionRegistry)
+        val header = JsKlibMetadataProtoBuf.Header.parseDelimitedFrom(inputStream, JsKlibMetadataSerializerProtocol.extensionRegistry)
+        val content = JsKlibMetadataProtoBuf.Library.parseFrom(inputStream, JsKlibMetadataSerializerProtocol.extensionRegistry)
         return JsKlibMetadataParts(header, content.packageFragmentList, content.importedModuleList)
     }
 }