[LL API] Pass a contextual module to ProjectStructureProvider

In certain cases, it's impossible to determine which module owns a
particular file without knowing the analysis context. For instance,
the file might be a part of a physical module, and be also included into
a virtual ad-hoc module (to be analyzed in separate, e.g. a VCS diff).

The new API allows to pass a contextual module. Basically it means
"give me a module for this element, implying that we are now analyzing a
contextual module".

^KT-57559 Fixed

analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleProvider.kt

@@ -12,11 +12,13 @@ import org.jetbrains.kotlin.serialization.deserialization.builtins.BuiltInSerial
 
 class KtStaticModuleProvider(
     private val builtinsModule: KtBuiltinsModule,
-    val projectStructure: KtModuleProjectStructure,
+    private val projectStructure: KtModuleProjectStructure,
 ) : ProjectStructureProvider() {
+    val allModules: List<KtModule>
+        get() = projectStructure.allKtModules()
 
     @OptIn(KtModuleStructureInternals::class)
-    override fun getKtModuleForKtElement(element: PsiElement): KtModule {
+    override fun getModule(element: PsiElement, contextualModule: KtModule?): KtModule {
         val containingFileAsPsiFile = element.containingFile
         val containingFileAsVirtualFile = containingFileAsPsiFile.virtualFile
         if (containingFileAsVirtualFile.extension == BuiltInSerializerProtocol.BUILTINS_FILE_EXTENSION) {
@@ -30,8 +32,7 @@ class KtStaticModuleProvider(
         }?.let { return it }
 
         return projectStructure.mainModules
-            .first { module ->
-                element in module.ktModule.contentScope
-            }.ktModule
+            .first { module -> element in module.ktModule.contentScope }
+            .ktModule
     }
 }

analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/project/structure/impl/KtModuleProviderImpl.kt

@@ -36,7 +36,7 @@ internal class KtModuleProviderImpl(
         LLFirBuiltinsSessionFactory.getInstance(project).getBuiltinsSession(platform).ktModule as KtBuiltinsModule
     }
 
-    override fun getKtModuleForKtElement(element: PsiElement): KtModule {
+    override fun getModule(element: PsiElement, contextualModule: KtModule?): KtModule {
         val containingFileAsPsiFile = element.containingFile
             ?: return ktNotUnderContentRootModuleWithoutPsiFile
         // If an [element] is created on the fly, e.g., via [KtPsiFactory],

analysis/analysis-api-standalone/tests/org/jetbrains/kotlin/analysis/api/standalone/fir/test/StandaloneModeTestServiceRegistrar.kt

@@ -31,7 +31,7 @@ public object StandaloneModeTestServiceRegistrar : AnalysisApiTestServiceRegistr
     override fun registerProjectModelServices(project: MockProject, testServices: TestServices) {
         val projectStructureProvider = project.getService(ProjectStructureProvider::class.java)
         val binaryModules =
-            (projectStructureProvider as? KtStaticModuleProvider)?.projectStructure?.allKtModules()?.filterIsInstance<KtBinaryModule>()
+            (projectStructureProvider as? KtStaticModuleProvider)?.allModules?.filterIsInstance<KtBinaryModule>()
                 ?: emptyList()
         val projectEnvironment = testServices.environmentManager.getProjectEnvironment()
         project.apply {