[PLUGIN API] Implement custom linkage for plugin extensions

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendFacade.kt

@@ -84,7 +84,7 @@ object JvmBackendFacade {
 
         stubGenerator.setIrProviders(irProviders)
 
-        val irModuleFragment = psi2ir.generateModuleFragment(psi2irContext, files, irProviders, expectDescriptorToSymbol = null)
+        val irModuleFragment = psi2ir.generateModuleFragment(psi2irContext, files, irProviders, pluginExtensions, expectDescriptorToSymbol = null)
         irLinker.postProcess()
 
         stubGenerator.unboundSymbolGeneration = true

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/Psi2IrTranslator.kt

@@ -58,7 +58,7 @@ class Psi2IrTranslator(
         val irProviders = generateTypicalIrProviderList(
             moduleDescriptor, context.irBuiltIns, context.symbolTable, extensions = generatorExtensions
         )
-        return generateModuleFragment(context, ktFiles, irProviders)
+        return generateModuleFragment(context, ktFiles, irProviders, emptyList())
     }
 
     fun createGeneratorContext(
@@ -76,6 +76,7 @@ class Psi2IrTranslator(
         context: GeneratorContext,
         ktFiles: Collection<KtFile>,
         irProviders: List<IrProvider>,
+        linkerExtensions: Collection<IrDeserializer.IrLinkerExtension>,
         expectDescriptorToSymbol: MutableMap<DeclarationDescriptor, IrSymbol>? = null
     ): IrModuleFragment {
         val moduleGenerator = ModuleGenerator(context)
@@ -86,7 +87,7 @@ class Psi2IrTranslator(
         postprocess(context, irModule)
 
         val deserializers = irProviders.filterIsInstance<IrDeserializer>()
-        deserializers.forEach { it.init(irModule) }
+        deserializers.forEach { it.init(irModule, linkerExtensions) }
 
         moduleGenerator.generateUnboundSymbolsAsDependencies(irProviders)
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/linkage/IrDeserializer.kt

@@ -15,6 +15,6 @@ interface IrDeserializer : IrProvider {
         fun resolveSymbol(symbol: IrSymbol): IrDeclaration? = null
     }
 
-    fun init(moduleFragment: IrModuleFragment?) {}
+    fun init(moduleFragment: IrModuleFragment?, extensions: Collection<IrLinkerExtension>) {}
     fun postProcess() {}
 }

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/KotlinIrLinker.kt

@@ -36,6 +36,7 @@ import org.jetbrains.kotlin.protobuf.CodedInputStream
 import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite.newInstance
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 import org.jetbrains.kotlin.backend.common.serialization.proto.Actual as ProtoActual
 import org.jetbrains.kotlin.backend.common.serialization.proto.IdSignature as ProtoIdSignature
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrConstructorCall as ProtoConstructorCall
@@ -67,6 +68,8 @@ abstract class KotlinIrLinker(
 
     private val haveSeen = mutableSetOf<IrSymbol>()
 
+    private lateinit var linkerExtensions: Collection<IrDeserializer.IrLinkerExtension>
+
     abstract inner class BasicIrModuleDeserializer(moduleDescriptor: ModuleDescriptor, override val klib: IrLibrary, override val strategy: DeserializationStrategy) :
         IrModuleDeserializer(moduleDescriptor) {
 
@@ -518,6 +521,22 @@ abstract class KotlinIrLinker(
 
     protected open fun platformSpecificSymbol(symbol: IrSymbol): Boolean = false
 
+    private fun tryResolveCustomDeclaration(symbol: IrSymbol): IrDeclaration? {
+        val descriptor = symbol.descriptor
+
+        if (descriptor is WrappedDeclarationDescriptor<*>) return null
+        if (descriptor is CallableMemberDescriptor) {
+            if (descriptor.kind == CallableMemberDescriptor.Kind.FAKE_OVERRIDE) {
+                // skip fake overrides
+                return null
+            }
+        }
+
+        return linkerExtensions.firstNotNullResult { it.resolveSymbol(symbol) }?.also {
+            require(symbol.owner == it)
+        }
+    }
+
     override fun getDeclaration(symbol: IrSymbol): IrDeclaration? {
 
         if (!symbol.isPublicApi) {
@@ -529,7 +548,7 @@ abstract class KotlinIrLinker(
         }
 
         if (!symbol.isBound) {
-            findDeserializedDeclarationForSymbol(symbol) ?: return null
+            findDeserializedDeclarationForSymbol(symbol) ?: tryResolveCustomDeclaration(symbol) ?: return null
         }
 
         // TODO: we do have serializations for those, but let's just create a stub for now.
@@ -548,7 +567,8 @@ abstract class KotlinIrLinker(
     protected open fun createCurrentModuleDeserializer(moduleFragment: IrModuleFragment, dependencies: Collection<IrModuleDeserializer>): IrModuleDeserializer =
         CurrentModuleDeserializer(moduleFragment, dependencies)
 
-    override fun init(moduleFragment: IrModuleFragment?) {
+    override fun init(moduleFragment: IrModuleFragment?, extensions: Collection<IrDeserializer.IrLinkerExtension>) {
+        linkerExtensions = extensions
         if (moduleFragment != null) {
             val currentModuleDependencies = moduleFragment.descriptor.allDependencyModules.map {
                 deserializersForModules[it] ?: error("No deserializer found for $it")

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

@@ -273,7 +273,7 @@ fun loadIr(
 
             val moduleFragment = deserializedModuleFragments.last()
 
-            irLinker.init(null)
+            irLinker.init(null, emptyList())
             ExternalDependenciesGenerator(symbolTable, listOf(irLinker), configuration.languageVersionSettings).generateUnboundSymbolsAsDependencies()
             irLinker.postProcess()
 
@@ -325,7 +325,7 @@ fun GeneratorContext.generateModuleFragmentWithPlugins(
         }
     }
 
-    return psi2Ir.generateModuleFragment(this, files, listOf(irLinker), expectDescriptorToSymbol)
+    return psi2Ir.generateModuleFragment(this, files, listOf(irLinker), extensions, expectDescriptorToSymbol)
 }
 
 private fun createBuiltIns(storageManager: StorageManager) = object : KotlinBuiltIns(storageManager) {}

js/js.tests/test/org/jetbrains/kotlin/benchmarks/GenerateIrRuntime.kt

@@ -478,7 +478,7 @@ class GenerateIrRuntime {
         val irProviders = listOf(irLinker)
 
         val psi2IrTranslator = Psi2IrTranslator(languageVersionSettings, psi2IrContext.configuration, signaturer)
-        return psi2IrTranslator.generateModuleFragment(psi2IrContext, files, irProviders, null)
+        return psi2IrTranslator.generateModuleFragment(psi2IrContext, files, irProviders, emptyList(), null)
     }
 
     private fun doSerializeModule(moduleFragment: IrModuleFragment, bindingContext: BindingContext, files: List<KtFile>, perFile: Boolean = false): String {
@@ -531,7 +531,7 @@ class GenerateIrRuntime {
         val jsLinker = JsIrLinker(moduleDescriptor, logger, irBuiltIns, symbolTable, functionFactory, null)
 
         val moduleFragment = jsLinker.deserializeFullModule(moduleDescriptor, moduleDescriptor.kotlinLibrary)
-        jsLinker.init(null)
+        jsLinker.init(null, emptyList())
         // Create stubs
         ExternalDependenciesGenerator(symbolTable, listOf(jsLinker), languageVersionSettings)
             .generateUnboundSymbolsAsDependencies()
@@ -562,7 +562,7 @@ class GenerateIrRuntime {
 
         val moduleFragment = jsLinker.deserializeFullModule(moduleDescriptor, moduleDescriptor.kotlinLibrary)
         // Create stubs
-        jsLinker.init(null)
+        jsLinker.init(null, emptyList())
         // Create stubs
         ExternalDependenciesGenerator(symbolTable, listOf(jsLinker), languageVersionSettings)
             .generateUnboundSymbolsAsDependencies()

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/repl/js/JsCoreScriptingCompiler.kt

@@ -65,7 +65,7 @@ class JsCoreScriptingCompiler(
         val psi2ir = Psi2IrTranslator(environment.configuration.languageVersionSettings, signaturer = signaturer)
         val psi2irContext = psi2ir.createGeneratorContext(module, bindingContext, symbolTable = symbolTable)
         val providers = generateTypicalIrProviderList(module, psi2irContext.irBuiltIns, psi2irContext.symbolTable)
-        val irModuleFragment = psi2ir.generateModuleFragment(psi2irContext, files, providers, null) // TODO: deserializer
+        val irModuleFragment = psi2ir.generateModuleFragment(psi2irContext, files, providers, emptyList(), null) // TODO: deserializer
 
         val context = JsIrBackendContext(
             irModuleFragment.descriptor,

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/repl/js/JsScriptDependencyCompiler.kt

@@ -56,7 +56,7 @@ class JsScriptDependencyCompiler(
         val moduleFragment = irDependencies.last()
         val irProviders = listOf(jsLinker)
 
-        jsLinker.init(null)
+        jsLinker.init(null, emptyList())
 
         ExternalDependenciesGenerator(symbolTable, irProviders, configuration.languageVersionSettings)
             .generateUnboundSymbolsAsDependencies()