From 79f1779d89e9997c31b139a2df89a2b773ee7c33 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 30 Sep 2021 16:12:54 +0300
Subject: [PATCH] [FIR] Add methods to collect all class and callable ids from
 extensions

They will be needed for generating IR for all declarations from plugins
---
 .../FirDeclarationGenerationExtension.kt           |  3 +++
 .../fir/plugin/generators/AllOpenClassGenerator.kt |  4 ++++
 .../AllOpenTopLevelDeclarationsGenerator.kt        | 14 +++++++++++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt
index b5ef21f3078..093fd9c595d 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt
@@ -38,8 +38,11 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirPredi
     // Can be called on IMPORTS stage
     open fun hasPackage(packageFqName: FqName): Boolean = false
 
+    // Can be called after BODY_RESOLVE stage (checkers and fir2ir)
     open fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> = emptySet()
     open fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> = emptySet()
+    open fun getTopLevelCallableIds(): Set<CallableId> = emptySet()
+    open fun getTopLevelClassIds(): Set<ClassId> = emptySet()
 
     fun interface Factory : FirExtension.Factory<FirDeclarationGenerationExtension>
 }
diff --git a/plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenClassGenerator.kt b/plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenClassGenerator.kt
index 4123665acb0..7f3ef37144d 100644
--- a/plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenClassGenerator.kt
+++ b/plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenClassGenerator.kt
@@ -151,6 +151,10 @@ class AllOpenClassGenerator(session: FirSession) : FirDeclarationGenerationExten
         }
     }
 
+    override fun getTopLevelClassIds(): Set<ClassId> {
+        return setOf(GENERATED_CLASS_ID)
+    }
+
     override fun hasPackage(packageFqName: FqName): Boolean {
         return packageFqName == FOO_PACKAGE
     }
diff --git a/plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenTopLevelDeclarationsGenerator.kt b/plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenTopLevelDeclarationsGenerator.kt
index d52f4174c4b..1e089434cfa 100644
--- a/plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenTopLevelDeclarationsGenerator.kt
+++ b/plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AllOpenTopLevelDeclarationsGenerator.kt
@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
 class AllOpenTopLevelDeclarationsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
@@ -66,7 +67,18 @@ class AllOpenTopLevelDeclarationsGenerator(session: FirSession) : FirDeclaration
         if (callableId.classId != null) return null
         return matchedClasses
             .filter { it.classId.packageFqName == callableId.packageName }
-            .firstOrNull { callableId.callableName.identifier == "dummy${it.classId.shortClassName.identifier}" }
+            .firstOrNull { callableId.callableName.identifier == it.classId.toDummyCallableName() }
+    }
+
+    private fun ClassId.toDummyCallableName(): String {
+        return "dummy${shortClassName.identifier}"
+    }
+
+    override fun getTopLevelCallableIds(): Set<CallableId> {
+        return matchedClasses.mapTo(mutableSetOf()) {
+            val classId = it.classId
+            CallableId(classId.packageFqName, Name.identifier(classId.toDummyCallableName()))
+        }
     }
 
     override val key: AllOpenPluginKey