From 87b9ce0a549fdefea3dd8c9058add9d52a6ec4d3 Mon Sep 17 00:00:00 2001
From: vladislavgrecko <vladislav.grechko@jetbrains.com>
Date: Wed, 4 Dec 2024 16:03:55 +0100
Subject: [PATCH] [JVM] Retrieve module names using experimental KA API for
 code fragments

At the moment, it's the only way to get actual module name for the
declaration when compiling code fragments from the IDEA debugger.
See KT-69226.

^KT-72500: Fixed
---
 .../analysis/api/projectStructure/KaModule.kt  | 18 +++++++++---------
 .../fir/projectStructure/LLFirModuleData.kt    |  3 +++
 .../kotlin/fir/lazy/Fir2IrLazyClass.kt         |  9 ++++++++-
 .../org/jetbrains/kotlin/fir/FirModuleData.kt  | 11 +++++++++++
 4 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/projectStructure/KaModule.kt b/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/projectStructure/KaModule.kt
index e3ef0b832236d..f39119fdeb5e1 100644
--- a/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/projectStructure/KaModule.kt
+++ b/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/projectStructure/KaModule.kt
@@ -83,15 +83,6 @@ public interface KaModule {
      */
     @KaExperimentalApi
     public val moduleDescription: String
-}
-
-/**
- * A module which consists of a set of source declarations inside a project.
- *
- * Generally, a main or test Source Set.
- */
-public interface KaSourceModule : KaModule {
-    public val name: String
 
     /**
      * A stable binary name of module from the *Kotlin* point of view.
@@ -102,6 +93,15 @@ public interface KaSourceModule : KaModule {
     @KaExperimentalApi
     public val stableModuleName: String?
         get() = null
+}
+
+/**
+ * A module which consists of a set of source declarations inside a project.
+ *
+ * Generally, a main or test Source Set.
+ */
+public interface KaSourceModule : KaModule {
+    public val name: String
 
     @KaExperimentalApi
     override val moduleDescription: String
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/projectStructure/LLFirModuleData.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/projectStructure/LLFirModuleData.kt
index eb88491c52983..66497b99ef706 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/projectStructure/LLFirModuleData.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/projectStructure/LLFirModuleData.kt
@@ -61,6 +61,9 @@ class LLFirModuleData private constructor(val ktModule: KaModule) : FirModuleDat
         get() = boundSession?.let { it as LLFirSession }
             ?: LLFirSessionCache.getInstance(ktModule.project).getSession(ktModule, preferBinary = true)
 
+    override val stableModuleName: String?
+        get() = ktModule.stableModuleName
+
     override fun equals(other: Any?): Boolean = this === other || other is LLFirModuleData && ktModule == other.ktModule
     override fun hashCode(): Int = ktModule.hashCode()
 }
diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt
index 5c59295aad898..48f7a6df32244 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt
+++ b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt
@@ -258,7 +258,14 @@ class Fir2IrLazyClass(
         set(_) = error("We should never need to store metadata of external declarations.")
 
     override val moduleName: String?
-        get() = fir.moduleName
+        get() {
+            fir.moduleName?.let { return it }
+            val moduleNameFromModuleData = fir.moduleData.stableModuleName ?: return null
+            require(moduleNameFromModuleData.startsWith("<") && moduleNameFromModuleData.endsWith(">")) {
+                "Stable module name is expected to be wrapped in '<>' brackets, but got `$moduleNameFromModuleData` instead"
+            }
+            return moduleNameFromModuleData.substring(1, moduleNameFromModuleData.length - 1)
+        }
 
     override val isNewPlaceForBodyGeneration: Boolean
         get() = fir.isNewPlaceForBodyGeneration == true
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirModuleData.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirModuleData.kt
index 8a77d1c3ade7b..fccd23ec61210 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirModuleData.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirModuleData.kt
@@ -72,6 +72,14 @@ abstract class FirModuleData : FirSessionComponent {
     override fun toString(): String {
         return "Module $name"
     }
+
+    /**
+     * For the Analysis API-originated modules, returns the stable module name of the corresponding KaModule.
+     * See [org.jetbrains.kotlin.analysis.api.projectStructure.KaModule.stableModuleName]
+     *
+     * For others, returns null.
+     */
+    abstract val stableModuleName: String?
 }
 
 class FirModuleDataImpl(
@@ -87,6 +95,9 @@ class FirModuleDataImpl(
         get() = boundSession
             ?: error("module data ${this::class.simpleName}:${name} not bound to session")
 
+    override val stableModuleName: String?
+        get() = null
+
     override val allDependsOnDependencies: List<FirModuleData> = topologicalSort(dependsOnDependencies) { it.dependsOnDependencies }
 }