From 22dfe07580af55da6a736b8763560bb9ad7c7feb Mon Sep 17 00:00:00 2001
From: Alexander Korepanov <alexander.korepanov@jetbrains.com>
Date: Mon, 8 May 2023 16:40:15 +0200
Subject: [PATCH] [IR BE] Add the ability to collect PL stubbed symbols

^KT-57347 related
---
 .../linkage/partial/MissingDeclarationStubGenerator.kt     | 6 ++++++
 .../linkage/partial/PartialLinkageSupportForLinker.kt      | 7 +++++++
 .../linkage/partial/PartialLinkageSupportForLinkerImpl.kt  | 5 +++++
 3 files changed, 18 insertions(+)

diff --git a/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/MissingDeclarationStubGenerator.kt b/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/MissingDeclarationStubGenerator.kt
index 99d9bd45bef..f2df00ef6de 100644
--- a/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/MissingDeclarationStubGenerator.kt
+++ b/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/MissingDeclarationStubGenerator.kt
@@ -38,6 +38,10 @@ internal class MissingDeclarationStubGenerator(private val builtIns: IrBuiltIns)
 
     private var declarationsToPatch = arrayListOf<IrDeclaration>()
 
+    private val stubbedSymbols = hashSetOf<IrSymbol>()
+
+    val allStubbedSymbols: Set<IrSymbol> get() = stubbedSymbols
+
     fun grabDeclarationsToPatch(): Collection<IrDeclaration> {
         val result = declarationsToPatch
         declarationsToPatch = arrayListOf()
@@ -47,6 +51,8 @@ internal class MissingDeclarationStubGenerator(private val builtIns: IrBuiltIns)
     override fun getDeclaration(symbol: IrSymbol): IrDeclaration {
         require(!symbol.isBound)
 
+        stubbedSymbols.add(symbol)
+
         return when (symbol) {
             is IrClassSymbol -> generateClass(symbol)
             is IrSimpleFunctionSymbol -> generateSimpleFunction(symbol)
diff --git a/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/PartialLinkageSupportForLinker.kt b/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/PartialLinkageSupportForLinker.kt
index 0e8bec99c1f..0f57fa87cec 100644
--- a/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/PartialLinkageSupportForLinker.kt
+++ b/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/PartialLinkageSupportForLinker.kt
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.backend.common.overrides.FakeOverrideBuilder
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.util.SymbolTable
 
 interface PartialLinkageSupportForLinker {
@@ -48,6 +49,11 @@ interface PartialLinkageSupportForLinker {
     fun generateStubsAndPatchUsages(symbolTable: SymbolTable, roots: () -> Sequence<IrModuleFragment>)
     fun generateStubsAndPatchUsages(symbolTable: SymbolTable, root: IrDeclaration)
 
+    /**
+     * Collect all symbols which were stubbed
+     */
+    fun collectAllStubbedSymbols(): Set<IrSymbol>
+
     companion object {
         val DISABLED = object : PartialLinkageSupportForLinker {
             override val isEnabled get() = false
@@ -56,6 +62,7 @@ interface PartialLinkageSupportForLinker {
             override fun exploreClassifiersInInlineLazyIrFunction(function: IrFunction) = Unit
             override fun generateStubsAndPatchUsages(symbolTable: SymbolTable, roots: () -> Sequence<IrModuleFragment>) = Unit
             override fun generateStubsAndPatchUsages(symbolTable: SymbolTable, root: IrDeclaration) = Unit
+            override fun collectAllStubbedSymbols(): Set<IrSymbol> = emptySet()
         }
     }
 }
diff --git a/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/PartialLinkageSupportForLinkerImpl.kt b/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/PartialLinkageSupportForLinkerImpl.kt
index 5f73b9d898a..b92e4f15d8d 100644
--- a/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/PartialLinkageSupportForLinkerImpl.kt
+++ b/compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/linkage/partial/PartialLinkageSupportForLinkerImpl.kt
@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.linkage.partial.PartialLinkageConfig
 import org.jetbrains.kotlin.ir.linkage.partial.PartialLinkageLogger
+import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.util.IrMessageLogger
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.allUnbound
@@ -66,6 +67,10 @@ internal class PartialLinkageSupportForLinkerImpl(
         generateStubsAndPatchUsagesInternal(symbolTable) { patcher.patchDeclarations(listOf(root)) }
     }
 
+    override fun collectAllStubbedSymbols(): Set<IrSymbol> {
+        return stubGenerator.allStubbedSymbols
+    }
+
     private fun generateStubsAndPatchUsagesInternal(symbolTable: SymbolTable, patchIrTree: () -> Unit) {
         // Generate stubs.
         for (symbol in symbolTable.allUnbound) {