[analysis api] provide a function to get containing module for a KtSymbol

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolContainingDeclarationProvider.kt

@@ -14,11 +14,13 @@ import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirReceiverParameterSymbol
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirTypeParameterSymbol
+import org.jetbrains.kotlin.analysis.api.fir.utils.getContainingKtModule
 import org.jetbrains.kotlin.analysis.api.impl.barebone.parentOfType
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolKind
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithKind
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.psi.*
 
@@ -62,6 +64,10 @@ internal class KtFirSymbolContainingDeclarationProvider(
         }
     }
 
+    override fun getContainingModule(symbol: KtSymbol): KtModule {
+       return symbol.getContainingKtModule(analysisSession.firResolveState)
+    }
+
     private fun getContainingDeclarationForBackingFieldSymbol(symbol: KtSymbolWithKind): KtSymbolWithKind {
         require(symbol is KtBackingFieldSymbol)
         return symbol.owningProperty

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/utils/ktSymbolUtils.kt

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.api.fir.utils
+
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirModuleResolveState
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.KtModuleBasedModuleData
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirModuleSession
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
+import org.jetbrains.kotlin.fir.dispatchReceiverClassOrNull
+import org.jetbrains.kotlin.fir.resolve.toFirRegularClass
+import org.jetbrains.kotlin.fir.resolve.toFirRegularClassSymbol
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+
+val KtSymbol.firSymbol: FirBasedSymbol<*>
+    get() {
+        require(this is KtFirSymbol<*>)
+        return this.firSymbol
+    }
+
+
+fun FirBasedSymbol<*>.getContainingKtModule(resolveState: LLFirModuleResolveState): KtModule {
+    val target = when (this) {
+        // callable fake overrides have use-site FirModuleData
+        is FirCallableSymbol -> dispatchReceiverClassOrNull()?.toFirRegularClassSymbol(resolveState.rootModuleSession)
+            ?: this
+
+        else -> this
+    }
+    val moduleData = target.moduleData
+    requireIsInstance<KtModuleBasedModuleData>(moduleData)
+    return moduleData.module
+}
+
+fun KtSymbol.getContainingKtModule(resolveState: LLFirModuleResolveState): KtModule =
+    firSymbol.getContainingKtModule(resolveState)

analysis/analysis-api/build.gradle.kts

@@ -17,6 +17,7 @@ dependencies {
     implementation(project(":analysis:analysis-internal-utils"))
     implementation(project(":analysis:analysis-api-providers"))
     implementation(project(":analysis:kt-references"))
+    api(project(":analysis:project-structure"))
 
     api(intellijCore())
     api(commonDependency("org.jetbrains.intellij.deps:asm-all"))

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtSymbolContainingDeclarationProvider.kt

@@ -7,9 +7,12 @@ package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithKind
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
 
 public abstract class KtSymbolContainingDeclarationProvider : KtAnalysisSessionComponent() {
     public abstract fun getContainingDeclaration(symbol: KtSymbol): KtSymbolWithKind?
+
+    public abstract fun getContainingModule(symbol: KtSymbol): KtModule
 }
 
 public interface KtSymbolContainingDeclarationProviderMixIn : KtAnalysisSessionMixIn {
@@ -21,4 +24,7 @@ public interface KtSymbolContainingDeclarationProviderMixIn : KtAnalysisSessionM
      */
     public fun KtSymbol.getContainingSymbol(): KtSymbolWithKind? =
         analysisSession.containingDeclarationProvider.getContainingDeclaration(this)
+
+    public fun KtSymbol.getContainingModule(): KtModule =
+        analysisSession.containingDeclarationProvider.getContainingModule(this)
 }

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/symbols/DebugSymbolRenderer.kt

@@ -11,11 +11,13 @@ import org.jetbrains.kotlin.analysis.api.KtConstantInitializerValue
 import org.jetbrains.kotlin.analysis.api.KtInitializerValue
 import org.jetbrains.kotlin.analysis.api.KtNonConstantInitializerValue
 import org.jetbrains.kotlin.analysis.api.annotations.*
+import org.jetbrains.kotlin.analysis.api.components.KtSymbolContainingDeclarationProviderMixIn
 import org.jetbrains.kotlin.analysis.api.components.KtSymbolInfoProviderMixIn
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtNamedSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtPossiblyNamedSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtClassErrorType
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.utils.printer.PrettyPrinter
 import org.jetbrains.kotlin.analysis.utils.printer.prettyPrint
 import org.jetbrains.kotlin.descriptors.Visibility
@@ -27,9 +29,9 @@ import org.jetbrains.kotlin.renderer.render
 import org.jetbrains.kotlin.resolve.deprecation.DeprecationInfo
 import java.lang.reflect.InvocationTargetException
 import kotlin.reflect.KClass
+import kotlin.reflect.KFunction
 import kotlin.reflect.KProperty
-import kotlin.reflect.full.declaredMemberExtensionProperties
+import kotlin.reflect.full.*
-import kotlin.reflect.full.extensionReceiverParameter
 
 public object DebugSymbolRenderer {
     public fun render(symbol: KtSymbol): String = prettyPrint { renderSymbol(symbol) }
@@ -49,6 +51,13 @@ public object DebugSymbolRenderer {
                 renderType(dispatchType)
             }
 
+            KtSymbolContainingDeclarationProviderMixIn::class
+                .declaredMemberExtensionFunctions
+                .filter { it.name == "getContainingModule" }
+                .forEach {
+                    renderFunction(it, this@renderExtra, symbol)
+                }
+
             KtSymbolInfoProviderMixIn::class.declaredMemberExtensionProperties
                 .asSequence()
                 .filter { (it.extensionReceiverParameter?.type?.classifier as? KClass<*>)?.isInstance(symbol) == true }
@@ -70,10 +79,19 @@ public object DebugSymbolRenderer {
         }
     }
 
-    private fun PrettyPrinter.renderProperty(property: KProperty<*>, vararg args: Any) {
+    private fun PrettyPrinter.renderFunction(function: KFunction<*>, vararg args: Any) {
+        appendLine().append(function.name).append(": ")
+        renderFunctionCall(function, args)
+    }
+
+    private fun PrettyPrinter.renderProperty(property: KProperty<*>, vararg args: Any ) {
+        appendLine().append(property.name).append(": ")
+        renderFunctionCall(property.getter, args)
+    }
+
+    private fun PrettyPrinter.renderFunctionCall(function: KFunction<*>, args: Array<out Any>) {
         try {
-            appendLine().append(property.name).append(": ")
+            renderValue(function.call(*args))
-            renderValue(property.getter.call(*args))
         } catch (e: InvocationTargetException) {
             append("Could not render due to ").appendLine(e.cause.toString())
         }
@@ -187,6 +205,7 @@ public object DebugSymbolRenderer {
             is KtInitializerValue -> renderKtInitializerValue(value)
             is KtAnnotationApplication -> renderAnnotationApplication(value)
             is KtAnnotationsList -> renderAnnotationsList(value)
+            is KtModule -> renderKtModule(value)
             // Other custom values
             is Name -> append(value.asString())
             is FqName -> append(value.asString())
@@ -205,6 +224,25 @@ public object DebugSymbolRenderer {
         }
     }
 
+    private fun PrettyPrinter.renderKtModule(ktModule: KtModule) {
+        val ktModuleClass = ktModule::class.allSuperclasses.filter { it in ktModuleSubclasses }.first()
+        append("${ktModuleClass.simpleName} \"${ktModule.moduleDescription}\"")
+    }
+
+    private fun KClass<*>.allSealedSubClasses(): List<KClass<*>> = buildList {
+        add(this@allSealedSubClasses)
+        sealedSubclasses.flatMapTo(this) { it.allSealedSubClasses() }
+    }
+
+    private val ktModuleSubclasses = KtModule::class.allSealedSubClasses().distinct().sortedWith { a, b ->
+        when {
+            a == b -> 0
+            a.isSubclassOf(b) -> -1
+            b.isSubclassOf(a) -> 1
+            else -> 0
+        }
+    }
+
     private fun PrettyPrinter.renderKtInitializerValue(value: KtInitializerValue) {
         when (value) {
             is KtConstantInitializerValue -> {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/KtModuleBasedModuleData.kt

@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
 
-internal class KtModuleBasedModuleData(
+class KtModuleBasedModuleData(
     val module: KtModule,
 ) : FirModuleData() {
     override val name: Name get() = Name.special("<${module.moduleDescription}>")