[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) {}