[IR] Provide new plugin API to access declaration via FqName in the safe way

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension.kt

@@ -8,15 +8,22 @@ package org.jetbrains.kotlin.backend.common.extensions
 import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.ir.BuiltinSymbolsBase
 import org.jetbrains.kotlin.config.LanguageVersionSettings
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.ir.builders.IrGeneratorContext
+import org.jetbrains.kotlin.ir.declarations.IrConstructor
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
+import org.jetbrains.kotlin.ir.symbols.*
+import org.jetbrains.kotlin.ir.util.IrDeserializer
 import org.jetbrains.kotlin.ir.util.IrExtensionGenerator
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.TypeTranslator
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.scopes.MemberScope
 
 class IrPluginContext(
     val moduleDescriptor: ModuleDescriptor,
@@ -25,8 +32,77 @@ class IrPluginContext(
     val symbolTable: SymbolTable,
     val typeTranslator: TypeTranslator,
     override val irBuiltIns: IrBuiltIns,
+    private val linker: IrDeserializer,
     val symbols: BuiltinSymbolsBase = BuiltinSymbolsBase(irBuiltIns, irBuiltIns.builtIns, symbolTable)
-) : IrGeneratorContext()
+) : IrGeneratorContext() {
+    private fun resolveMemberScope(fqName: FqName): MemberScope? {
+        val pkg = moduleDescriptor.getPackage(fqName)
+
+        if (fqName.isRoot || pkg.fragments.isNotEmpty()) return pkg.memberScope
+
+        val parentMemberScope = resolveMemberScope(fqName.parent()) ?: return null
+
+        val classDescriptor =
+            parentMemberScope.getContributedClassifier(fqName.shortName(), NoLookupLocation.FROM_BACKEND) as? ClassDescriptor ?: return null
+
+        return classDescriptor.unsubstitutedMemberScope
+    }
+
+    private fun <S : IrSymbol> resolveSymbol(fqName: FqName, referencer: (MemberScope) -> S?): S? {
+        val memberScope = resolveMemberScope(fqName) ?: return null
+
+        val symbol = referencer(memberScope) ?: return null
+        if (symbol.isBound) return symbol
+
+        linker.getDeclaration(symbol)
+        // TODO: post process
+//        linker.
+        return symbol
+    }
+
+    private fun <S : IrSymbol> resolveSymbolCollection(fqName: FqName, referencer: (MemberScope) -> Collection<S>): Collection<S> {
+        val memberScope = resolveMemberScope(fqName) ?: return emptyList()
+
+        val symbols = referencer(memberScope)
+
+        symbols.forEach { if (!it.isBound) linker.getDeclaration(it) }
+
+        // TODO: post process
+//        linker.
+        return symbols
+    }
+
+    fun referenceClass(fqName: FqName): IrClassSymbol? {
+        assert(!fqName.isRoot)
+        return resolveSymbol(fqName.parent()) { scope ->
+            val classDescriptor = scope.getContributedClassifier(fqName.shortName(), NoLookupLocation.FROM_BACKEND) as? ClassDescriptor?
+            classDescriptor?.let {
+                symbolTable.referenceClass(it)
+            }
+        }
+    }
+
+    fun referenceConstructors(classFqn: FqName): Collection<IrConstructorSymbol> {
+        val classSymbol = referenceClass(classFqn) ?: error("Cannot find class $classFqn")
+        return classSymbol.owner.declarations.filterIsInstance<IrConstructor>().map { it.symbol }
+    }
+
+    fun referenceFunctions(fqName: FqName): Collection<IrSimpleFunctionSymbol> {
+        assert(!fqName.isRoot)
+        return resolveSymbolCollection(fqName.parent()) { scope ->
+            val descriptors = scope.getContributedFunctions(fqName.shortName(), NoLookupLocation.FROM_BACKEND)
+            descriptors.map { symbolTable.referenceSimpleFunction(it) }
+        }
+    }
+
+    fun referenceProperties(fqName: FqName): Collection<IrPropertySymbol> {
+        assert(!fqName.isRoot)
+        return resolveSymbolCollection(fqName.parent()) { scope ->
+            val descriptors = scope.getContributedVariables(fqName.shortName(), NoLookupLocation.FROM_BACKEND)
+            descriptors.map { symbolTable.referenceProperty(it) }
+        }
+    }
+}
 
 interface IrGenerationExtension : IrExtensionGenerator {
     companion object :

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

@@ -40,22 +40,6 @@ object JvmBackendFacade {
         val functionFactory = IrFunctionFactory(psi2irContext.irBuiltIns, psi2irContext.symbolTable)
         psi2irContext.irBuiltIns.functionFactory = functionFactory
 
-        for (extension in pluginExtensions) {
-            psi2ir.addPostprocessingStep { module ->
-                extension.generate(
-                    module,
-                    IrPluginContext(
-                        psi2irContext.moduleDescriptor,
-                        psi2irContext.bindingContext,
-                        psi2irContext.languageVersionSettings,
-                        psi2irContext.symbolTable,
-                        psi2irContext.typeTranslator,
-                        psi2irContext.irBuiltIns
-                    )
-                )
-            }
-        }
-
         val stubGenerator = DeclarationStubGenerator(
             psi2irContext.moduleDescriptor, psi2irContext.symbolTable, psi2irContext.irBuiltIns.languageVersionSettings, extensions
         )
@@ -68,6 +52,24 @@ object JvmBackendFacade {
             stubGenerator,
             mangler
         )
+
+        for (extension in pluginExtensions) {
+            psi2ir.addPostprocessingStep { module ->
+                extension.generate(
+                    module,
+                    IrPluginContext(
+                        psi2irContext.moduleDescriptor,
+                        psi2irContext.bindingContext,
+                        psi2irContext.languageVersionSettings,
+                        psi2irContext.symbolTable,
+                        psi2irContext.typeTranslator,
+                        psi2irContext.irBuiltIns,
+                        linker = irLinker
+                    )
+                )
+            }
+        }
+
         val dependencies = psi2irContext.moduleDescriptor.allDependencyModules.map {
             val kotlinLibrary = (it.getCapability(KlibModuleOrigin.CAPABILITY) as? DeserializedKlibModuleOrigin)?.library
             irLinker.deserializeIrModuleHeader(it, kotlinLibrary)

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

@@ -304,7 +304,8 @@ fun GeneratorContext.generateModuleFragmentWithPlugins(
                     languageVersionSettings,
                     symbolTable,
                     typeTranslator,
-                    irBuiltIns
+                    irBuiltIns,
+                    linker = irLinker
                 )
             )
         }