[JS IR] Track PL stubs in IC infra

^KT-57347 fixed
2023-05-04 14:59:28 +02:00
parent 22dfe07580
commit fab8a101bb
5 changed files with 151 additions and 12 deletions
@@ -264,6 +264,17 @@ class CacheUpdater(
            return exportedSymbols
        }

+        fun collectStubbedSignatures(): Set<IdSignature> {
+            val stubbedSignatures = hashSetOf<IdSignature>()
+            for (cache in incrementalCaches.values) {
+                val fileStubbedSignatures = cache.collectFilesWithStubbedSignatures()
+                for (signatures in fileStubbedSignatures.values) {
+                    stubbedSignatures += signatures
+                }
+            }
+            return stubbedSignatures
+        }
+
        private fun KotlinSourceFileMutableMap<DirtyFileMetadata>.getExportedSignaturesAndAddMetadata(
            symbolProviders: List<FileSignatureProvider>,
            libFile: KotlinLibraryFile,
@@ -508,6 +519,26 @@ class CacheUpdater(
            return libFilesToRebuild
        }

+        fun collectFilesWithUpdatedStubbedSymbols(dirtyFiles: KotlinSourceFileMap<*>): KotlinSourceFileMap<KotlinSourceFileExports> {
+            val libFiles = KotlinSourceFileMutableMap<KotlinSourceFileExports>()
+
+            for ((libFile, cache) in incrementalCaches.entries) {
+                val filesToRebuild by lazy(LazyThreadSafetyMode.NONE) { libFiles.getOrPutFiles(libFile) }
+                val fileStats by lazy(LazyThreadSafetyMode.NONE) { dirtyFileStats.getOrPutFiles(libFile) }
+                val alreadyDirtyFiles = dirtyFiles[libFile]?.keys ?: emptySet()
+                val filesWithStubbedSignatures = cache.collectFilesWithStubbedSignatures()
+
+                for ((srcFile, stubbedSignatures) in filesWithStubbedSignatures.entries) {
+                    if (srcFile !in alreadyDirtyFiles && stubbedSignatures.any { it in signatureHashCalculator }) {
+                        filesToRebuild[srcFile] = cache.fetchSourceFileFullMetadata(srcFile)
+                        fileStats.addDirtFileStat(srcFile, DirtyFileState.UPDATED_IMPORTS)
+                    }
+                }
+            }
+
+            return libFiles
+        }
+
        fun updateStdlibIntrinsicDependencies(
            loadedIr: LoadedJsIr,
            mainModule: IrModuleFragment,
@@ -574,13 +605,15 @@ class CacheUpdater(
                    icError("can not delete cache directory ${it.cacheDir.absolutePath}")
                }
            }
+
+            val stubbedSignatures = loadedIr.collectSymbolsReplacedWithStubs().mapNotNullTo(hashSetOf()) { it.signature }
            return libraryDependencies.keys.associate { library ->
                val libFile = KotlinLibraryFile(library)
                val incrementalCache = getLibIncrementalCache(libFile)
                val providers = loadedIr.getSignatureProvidersForLib(libFile)
                val signatureToIndexMapping = providers.associate { KotlinSourceFile(it.irFile) to it.getSignatureToIndexMapping() }

-                val cacheArtifact = incrementalCache.buildIncrementalCacheArtifact(signatureToIndexMapping)
+                val cacheArtifact = incrementalCache.buildAndCommitCacheArtifact(signatureToIndexMapping, stubbedSignatures)

                val libFragment = loadedIr.loadedFragments[libFile] ?: notFoundIcError("loaded fragment", libFile)
                val sourceFilesFromCache = cacheArtifact.getSourceFiles()
@@ -650,13 +683,15 @@ class CacheUpdater(

        stopwatch.startNext("Modified files - collecting exported signatures")
        val dirtyFileExports = updater.collectExportedSymbolsForDirtyFiles(modifiedFiles)
+        val stubbedSignatures = updater.collectStubbedSignatures()

        stopwatch.startNext("Modified files - loading and linking IR")
        val jsIrLinkerLoader = JsIrLinkerLoader(
            compilerConfiguration = compilerConfiguration,
            dependencyGraph = updater.libraryDependencies,
            mainModuleFriends = updater.mainModuleFriendLibraries,
-            irFactory = irFactory()
+            irFactory = irFactory(),
+            stubbedSignatures = stubbedSignatures
        )
        var loadedIr = jsIrLinkerLoader.loadIr(dirtyFileExports)

@@ -672,16 +707,20 @@ class CacheUpdater(

            stopwatch.startNext("Dependencies ($iterations) - collecting exported signatures for files with updated exports and imports")
            val filesToRebuild = updater.collectFilesToRebuildSignatures(filesWithModifiedExportsOrImports)
+            dirtyFileExports.copyFilesFrom(filesToRebuild)

-            if (filesToRebuild.isEmpty()) {
+            stopwatch.startNext("Dependencies ($iterations) - collecting files that contain updated stubbed symbols")
+            val filesWithUpdatedStubbedSymbolsToRebuild = updater.collectFilesWithUpdatedStubbedSymbols(dirtyFileExports)
+            dirtyFileExports.copyFilesFrom(filesWithUpdatedStubbedSymbolsToRebuild)
+
+            lastDirtyFiles = filesToRebuild.combineWith(filesWithUpdatedStubbedSymbolsToRebuild)
+
+            if (lastDirtyFiles.isEmpty()) {
                break
            }

-            lastDirtyFiles = filesToRebuild
-            dirtyFileExports.copyFilesFrom(filesToRebuild)
-
            stopwatch.startNext("Dependencies ($iterations) - loading and linking IR for files with modified exports and imports")
-            loadedIr = jsIrLinkerLoader.loadIr(filesToRebuild)
+            loadedIr = jsIrLinkerLoader.loadIr(lastDirtyFiles)
            iterations++
        }

@@ -771,7 +810,7 @@ fun rebuildCacheForDirtyFiles(

    val modifiedFiles = mapOf(libFile to dirtySrcFiles.associateWith { emptyMetadata })

-    val jsIrLoader = JsIrLinkerLoader(configuration, dependencyGraph, emptyList(), irFactory)
+    val jsIrLoader = JsIrLinkerLoader(configuration, dependencyGraph, emptyList(), irFactory, emptySet())
    val loadedIr = jsIrLoader.loadIr(KotlinSourceFileMap<KotlinSourceFileExports>(modifiedFiles), true)

    val currentIrModule = loadedIr.loadedFragments[libFile] ?: notFoundIcError("loaded fragment", libFile)
@@ -126,4 +126,8 @@ internal class IdSignatureHashCalculator(private val icHasher: ICHasher) {
        idSignatureHashes[signature] = signatureHash
        return signatureHash
    }
+
+    operator fun contains(signature: IdSignature): Boolean {
+        return signature in idSignatureSources
+    }
 }
@@ -16,16 +16,19 @@ import java.io.File
 internal class IncrementalCache(private val library: KotlinLibraryHeader, val cacheDir: File) {
    companion object {
        private const val CACHE_HEADER = "ic.header.bin"
+        private const val STUBBED_SYMBOLS = "ic.stubbed-symbols.bin"

        private const val BINARY_AST_SUFFIX = "ast.bin"
        private const val METADATA_SUFFIX = "metadata.bin"
    }

    private val cacheHeaderFile = File(cacheDir, CACHE_HEADER)
+    private val stubbedSymbolsFile = File(cacheDir, STUBBED_SYMBOLS)

    private var cacheHeaderShouldBeUpdated = false

-    private var removedSrcFiles: Collection<KotlinSourceFile> = emptyList()
+    private var removedSrcFiles: Set<KotlinSourceFile> = emptySet()
+    private var modifiedSrcFiles: Set<KotlinSourceFile> = emptySet()

    private val kotlinLibrarySourceFileMetadata = hashMapOf<KotlinSourceFile, KotlinSourceFileMetadata>()

@@ -37,6 +40,10 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
        }
    }

+    private val filesWithStubbedSignatures: Map<KotlinSourceFile, Set<IdSignature>> by lazy {
+        fetchFilesWithStubbedSymbols()
+    }
+
    val libraryFileFromHeader by lazy(LazyThreadSafetyMode.NONE) { cacheHeaderFromDisk?.libraryFile }

    private class CacheHeader(
@@ -82,7 +89,10 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
        return File(cacheDir, "${File(path).name}.$pathHash.$suffix")
    }

-    fun buildIncrementalCacheArtifact(signatureToIndexMapping: Map<KotlinSourceFile, Map<IdSignature, Int>>): IncrementalCacheArtifact {
+    fun buildAndCommitCacheArtifact(
+        signatureToIndexMapping: Map<KotlinSourceFile, Map<IdSignature, Int>>,
+        stubbedSignatures: Set<IdSignature>
+    ): IncrementalCacheArtifact {
        val klibSrcFiles = if (cacheHeaderShouldBeUpdated) {
            val newCacheHeader = CacheHeader(library)
            cacheHeaderFile.useCodedOutput { newCacheHeader.toProtoStream(this) }
@@ -96,9 +106,24 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
            removedFile.getCacheFile(METADATA_SUFFIX).delete()
        }

+        val updatedFilesWithStubbedSignatures = hashMapOf<KotlinSourceFile, Set<IdSignature>>()
+
        val fileArtifacts = klibSrcFiles.map { srcFile ->
-            commitSourceFileMetadata(srcFile, signatureToIndexMapping[srcFile] ?: emptyMap())
+            val signatureMapping = signatureToIndexMapping[srcFile] ?: emptyMap()
+            val artifact = commitSourceFileMetadata(srcFile, signatureMapping)
+
+            val fileStubbedSignatures = when (artifact) {
+                is SourceFileCacheArtifact.CommitMetadata -> signatureMapping.keys.filterTo(hashSetOf()) { it in stubbedSignatures }
+                else -> filesWithStubbedSignatures[srcFile] ?: emptySet()
+            }
+            if (fileStubbedSignatures.isNotEmpty()) {
+                updatedFilesWithStubbedSignatures[srcFile] = fileStubbedSignatures
+            }
+            artifact
        }
+
+        commitFilesWithStubbedSignatures(updatedFilesWithStubbedSignatures, signatureToIndexMapping)
+
        return IncrementalCacheArtifact(cacheDir, removedSrcFiles.isNotEmpty(), fileArtifacts, library.jsOutputName)
    }

@@ -132,6 +157,7 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
        }

        removedSrcFiles = removedFiles.keys
+        modifiedSrcFiles = modifiedFiles.keys
        cacheHeaderShouldBeUpdated = true

        return ModifiedFiles(addedFiles, removedFiles, modifiedFiles, nonModifiedFiles)
@@ -145,6 +171,55 @@ internal class IncrementalCache(private val library: KotlinLibraryHeader, val ca
        kotlinLibrarySourceFileMetadata[srcFile] = sourceFileMetadata
    }

+    fun collectFilesWithStubbedSignatures(): Map<KotlinSourceFile, Set<IdSignature>> {
+        return filesWithStubbedSignatures
+    }
+
+    private fun fetchFilesWithStubbedSymbols(): Map<KotlinSourceFile, Set<IdSignature>> {
+        return stubbedSymbolsFile.useCodedInputIfExists {
+            buildMapUntil(readInt32()) {
+                val srcFile = KotlinSourceFile.fromProtoStream(this@useCodedInputIfExists)
+                val signatureDeserializer = idSignatureSerialization.getIdSignatureDeserializer(srcFile)
+                if (srcFile in modifiedSrcFiles || srcFile in removedSrcFiles) {
+                    repeat(readInt32()) {
+                        signatureDeserializer.skipIdSignature(this@useCodedInputIfExists)
+                    }
+                } else {
+                    val unboundSignatures = buildSetUntil(readInt32()) {
+                        add(signatureDeserializer.deserializeIdSignature(this@useCodedInputIfExists))
+                    }
+                    put(srcFile, unboundSignatures)
+                }
+            }
+        } ?: emptyMap()
+    }
+
+    private fun commitFilesWithStubbedSignatures(
+        updatedFilesWithStubbedSignatures: Map<KotlinSourceFile, Set<IdSignature>>,
+        signatureToIndexMapping: Map<KotlinSourceFile, Map<IdSignature, Int>>,
+    ) {
+        if (updatedFilesWithStubbedSignatures.isEmpty()) {
+            stubbedSymbolsFile.delete()
+            return
+        }
+
+        if (updatedFilesWithStubbedSignatures == filesWithStubbedSignatures) {
+            return
+        }
+
+        stubbedSymbolsFile.useCodedOutput {
+            writeInt32NoTag(updatedFilesWithStubbedSignatures.size)
+            for ((srcFile, stubbedSignatures) in updatedFilesWithStubbedSignatures) {
+                val serializer = idSignatureSerialization.getIdSignatureSerializer(srcFile, signatureToIndexMapping[srcFile] ?: emptyMap())
+                srcFile.toProtoStream(this@useCodedOutput)
+                writeInt32NoTag(stubbedSignatures.size)
+                for (signature in stubbedSignatures) {
+                    serializer.serializeIdSignature(this@useCodedOutput, signature)
+                }
+            }
+        }
+    }
+
    private fun fetchSourceFileMetadata(srcFile: KotlinSourceFile, loadSignatures: Boolean) =
        kotlinLibrarySourceFileMetadata.getOrPut(srcFile) {
            val deserializer = idSignatureSerialization.getIdSignatureDeserializer(srcFile)
@@ -25,7 +25,9 @@ import org.jetbrains.kotlin.ir.declarations.IrFactory
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.descriptors.IrDescriptorBasedFunctionFactory
 import org.jetbrains.kotlin.ir.linkage.partial.partialLinkageConfig
+import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.util.ExternalDependenciesGenerator
+import org.jetbrains.kotlin.ir.util.IdSignature
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.irMessageLogger
 import org.jetbrains.kotlin.js.config.ErrorTolerancePolicy
@@ -75,13 +77,18 @@ internal data class LoadedJsIr(
        linker.checkNoUnboundSymbols(linker.symbolTable, "at the end of IR linkage process")
        linker.clear()
    }
+
+    fun collectSymbolsReplacedWithStubs(): Set<IrSymbol> {
+        return linker.partialLinkageSupport.collectAllStubbedSymbols()
+    }
 }

 internal class JsIrLinkerLoader(
    private val compilerConfiguration: CompilerConfiguration,
    private val dependencyGraph: Map<KotlinLibrary, List<KotlinLibrary>>,
    private val mainModuleFriends: Collection<KotlinLibrary>,
-    private val irFactory: IrFactory
+    private val irFactory: IrFactory,
+    private val stubbedSignatures: Set<IdSignature>
 ) {
    private val mainLibrary = dependencyGraph.keys.lastOrNull() ?: notFoundIcError("main library")

@@ -211,6 +218,12 @@ internal class JsIrLinkerLoader(
                        }
                    }
                }
+
+                for (stubbedSignature in stubbedSignatures) {
+                    if (stubbedSignature in moduleDeserializer) {
+                        moduleDeserializer.addModuleReachableTopLevel(stubbedSignature)
+                    }
+                }
            }
        }

@@ -84,6 +84,14 @@ fun <T> KotlinSourceFileMap<T>.toMutable(): KotlinSourceFileMutableMap<T> {
    return KotlinSourceFileMutableMap(entries.associateTo(HashMap(entries.size)) { it.key to HashMap(it.value) })
 }

+fun <T> KotlinSourceFileMap<T>.combineWith(other: KotlinSourceFileMap<T>): KotlinSourceFileMap<T> {
+    return when {
+        isEmpty() -> other
+        other.isEmpty() -> this
+        else -> toMutable().also { it.copyFilesFrom(other) }
+    }
+}
+
 fun KotlinSourceFileMap<Set<IdSignature>>.flatSignatures(): Set<IdSignature> {
    val allSignatures = hashSetOf<IdSignature>()
    forEachFile { _, _, signatures -> allSignatures += signatures }