[JS IR] Track PL stubs in IC infra

^KT-57347 fixed

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/ic/CacheUpdater.kt

@@ -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)

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/ic/IdSignatureHashCalculator.kt

@@ -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
+    }
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/ic/IncrementalCache.kt

@@ -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)

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/ic/JsIrLinkerLoader.kt

@@ -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)
+                    }
+                }
             }
         }
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/ic/KotlinSourceFileMetadata.kt

@@ -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 }