Extract module actualization from 'CompilationPeerCollector'

Always provide the JVM actualizer, so the 'CompilationPeerCollector'
can find JVM peers even for common code fragments.

The actualizer is deliberately made extensible to support KT-77426
in the future.

^KT-76457 Fixed

(cherry picked from commit a649676630b49c87b4c641554e16d7fb263d1070)
diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KaFirCompilerFacility.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KaFirCompilerFacility.kt
index b366b02..e0baad7 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KaFirCompilerFacility.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KaFirCompilerFacility.kt
@@ -29,7 +29,6 @@
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinCompilerPluginsProvider
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinModuleDependentsProvider
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinModuleOutputProvider
-import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinProjectStructureProvider
 import org.jetbrains.kotlin.analysis.api.projectStructure.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.*
@@ -39,6 +38,9 @@
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CompilationPeerData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.projectStructure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ImplementationPlatformKind
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLKindBasedPlatformActualizer
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLPlatformActualizer
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
@@ -192,9 +194,10 @@
             computeCodeFragmentMappings(mainFirFile, firResolveSession, configuration)
         }
 
-        val compilationPeerData = CompilationPeerCollector.process(mainFirFile)
+        val actualizer = LLKindBasedPlatformActualizer(ImplementationPlatformKind.JVM)
+        val compilationPeerData = CompilationPeerCollector.process(mainFirFile, actualizer)
 
-        val chunkRegistrar = CompilationChunkRegistrar(mainFile, mainFirFile, target)
+        val chunkRegistrar = CompilationChunkRegistrar(mainFile, mainFirFile, target, actualizer)
         val chunks = collectCompilationChunks(chunkRegistrar, compilationPeerData, codeFragmentMappings)
 
         val jvmIrDeserializer = JvmIrDeserializerImpl()
@@ -334,7 +337,8 @@
     private inner class CompilationChunkRegistrar(
         private val originalMainFile: KtFile,
         private val originalMainFirFile: FirFile,
-        private val target: KaCompilerTarget
+        private val target: KaCompilerTarget,
+        private val actualizer: LLPlatformActualizer?
     ) {
         private val originalMainModule = originalMainFirFile.llFirModuleData.ktModule
         private val originalMainContextModule = (originalMainModule as? KaDanglingFileModule)?.contextModule
@@ -416,11 +420,7 @@
             }
 
         private val KaModule.implementingJvmModule: KaModule?
-            get() {
-                return KotlinProjectStructureProvider.getInstance(project)
-                    .getImplementingModules(this)
-                    .find { it.targetPlatform.isJvm() }
-            }
+            get() = actualizer?.actualize(this)
 
         private val moduleCache = HashMap<KaModule, KaModule>()
 
diff --git a/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerated.java b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerated.java
index aab5cb2..32c7486 100644
--- a/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerated.java
+++ b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerated.java
@@ -290,6 +290,12 @@
     }
 
     @Test
+    @TestMetadata("commonOnlyUsageInlineTransitive4.kt")
+    public void testCommonOnlyUsageInlineTransitive4() {
+      runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.kt");
+    }
+
+    @Test
     @TestMetadata("constValFromLib.kt")
     public void testConstValFromLib() {
       runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/constValFromLib.kt");
diff --git a/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated.java b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated.java
index df1aed8..605e4a1 100644
--- a/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated.java
+++ b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated.java
@@ -290,6 +290,12 @@
     }
 
     @Test
+    @TestMetadata("commonOnlyUsageInlineTransitive4.kt")
+    public void testCommonOnlyUsageInlineTransitive4() {
+      runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.kt");
+    }
+
+    @Test
     @TestMetadata("constValFromLib.kt")
     public void testConstValFromLib() {
       runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/constValFromLib.kt");
diff --git a/analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.ir.txt b/analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.ir.txt
new file mode 100644
index 0000000..eac7820
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.ir.txt
@@ -0,0 +1,13 @@
+MODULE_FRAGMENT
+  FILE fqName:<root> fileName:fragment.kt
+    CLASS CLASS name:CodeFragment modality:FINAL visibility:public superTypes:[kotlin.Any]
+      thisReceiver: VALUE_PARAMETER INSTANCE_RECEIVER kind:DispatchReceiver name:<this> type:<root>.CodeFragment
+      CONSTRUCTOR visibility:public returnType:<root>.CodeFragment [primary]
+        BLOCK_BODY
+          DELEGATING_CONSTRUCTOR_CALL 'public constructor <init> () [primary,expect] declared in kotlin.Any'
+      FUN name:run visibility:public modality:FINAL returnType:kotlin.String
+        EXPRESSION_BODY
+          BLOCK type=kotlin.String origin=null
+            CALL 'public final fun plus (other: kotlin.Any?): kotlin.String [expect,operator] declared in kotlin.String' type=kotlin.String origin=PLUS
+              ARG <this>: CALL 'public final fun lib (): kotlin.String [inline] declared in test.JvmKt' type=kotlin.String origin=null
+              ARG other: CALL 'public final fun root (): kotlin.String [inline] declared in root.RootKt' type=kotlin.String origin=null
diff --git a/analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.kt b/analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.kt
new file mode 100644
index 0000000..d6e481e
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.kt
@@ -0,0 +1,50 @@
+// DUMP_IR
+// LANGUAGE: +MultiPlatformProjects
+
+// MODULE: common
+// TARGET_PLATFORM: Common
+
+// FILE: Common.kt
+package test
+
+expect inline fun lib(): String
+
+
+// MODULE: root
+// TARGET_PLATFORM: JVM
+
+// FILE: Root.kt
+package root
+
+inline fun root(): String = "root"
+
+
+// MODULE: jvm(root)()(common)
+// TARGET_PLATFORM: JVM
+
+// FILE: Jvm.kt
+package test
+import root.*
+
+actual inline fun lib(): String = root()
+
+
+// MODULE: app(common)
+// TARGET_PLATFORM: Common
+
+// FILE: Common.kt
+package app
+import test.*
+
+fun test() {
+    <caret_context>lib()
+}
+
+
+// MODULE: main
+// MODULE_KIND: CodeFragment
+// CONTEXT_MODULE: app
+
+// FILE: fragment.kt
+// CODE_FRAGMENT_KIND: EXPRESSION
+lib() + root.root()
\ No newline at end of file
diff --git a/analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.txt b/analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.txt
new file mode 100644
index 0000000..35849586
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/commonOnlyUsageInlineTransitive4.txt
@@ -0,0 +1,5 @@
+public final class CodeFragment {
+    // source: 'fragment.kt'
+    public method <init>(): void
+    public final static method run(): java.lang.String
+}
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/compile/CompilationPeerCollector.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/compile/CompilationPeerCollector.kt
index 72be926..9672ec7 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/compile/CompilationPeerCollector.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/compile/CompilationPeerCollector.kt
@@ -13,7 +13,7 @@
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaModule
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionService
 import org.jetbrains.kotlin.analysis.low.level.api.fir.projectStructure.llFirModuleData
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ImplementationPlatformKind
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLPlatformActualizer
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.containingKtFileIfAny
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
 import org.jetbrains.kotlin.codegen.state.GenerationState
@@ -49,12 +49,10 @@
  * Note that compiled declarations are not analyzed, as the backend can inline them natively.
  */
 @KaImplementationDetail
-class CompilationPeerCollector private constructor(private val platformKind: ImplementationPlatformKind?) {
+class CompilationPeerCollector private constructor(private val actualizer: LLPlatformActualizer?) {
     companion object {
-        fun process(file: FirFile): CompilationPeerData {
-            val platformKind = ImplementationPlatformKind.fromTargetPlatform(file.llFirModuleData.platform)
-
-            val collector = CompilationPeerCollector(platformKind)
+        fun process(file: FirFile, actualizer: LLPlatformActualizer?): CompilationPeerData {
+            val collector = CompilationPeerCollector(actualizer)
             collector.process(file)
 
             return CompilationPeerData(
@@ -98,7 +96,7 @@
         }
 
         // Avoid deep stacks by gathering callee files first
-        val visitor = CompilationPeerCollectingVisitor(ktFile.project, platformKind)
+        val visitor = CompilationPeerCollectingVisitor(ktFile.project, actualizer)
         file.accept(visitor)
 
         inlinedClasses.addAll(visitor.inlinedClasses)
@@ -150,7 +148,7 @@
 
 private class CompilationPeerCollectingVisitor(
     val project: Project,
-    val platformKind: ImplementationPlatformKind?,
+    val actualizer: LLPlatformActualizer?,
 ) : FirDefaultVisitorVoid() {
     private val collectedFunctions = HashSet<FirFunction>()
     private val collectedFiles = LinkedHashSet<FirFile>()
@@ -269,19 +267,23 @@
      * Instead, we need to compile the 'actual' counterpart that does have a body.
      */
     private fun registerActualCounterpart(originalFunction: FirFunction) {
-        if (platformKind == null) {
+        if (actualizer == null) {
             // We aren't sure which implementation platform to choose. So, aborting
             return
         }
 
         val originalPsi = originalFunction.source?.psi as? KtDeclaration ?: return
+        val originalModule = projectStructureProvider.getModule(originalPsi, useSiteModule = null)
+
+        val targetModule = actualizer.actualize(originalModule) ?: return
 
         // Across all 'actual' declarations, find those with a matching platform kind, and register their containing files
         for (actualPsi in actualDeclarationProvider?.getActualDeclarations(originalPsi).orEmpty()) {
             val actualPsiFile = actualPsi.containingFile as? KtFile ?: continue
-
             val actualModule = projectStructureProvider.getModule(actualPsiFile, useSiteModule = null)
-            if (platformKind.matches(actualModule.targetPlatform)) {
+
+            // The file we found is from the correct actualized module
+            if (targetModule == actualModule) {
                 val actualResolveSession = resolveSessionService.getFirResolveSession(actualModule)
                 val actualFile = actualResolveSession.getOrBuildFirFile(actualPsiFile)
                 collectedFiles.add(actualFile)
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ImplementationPlatformKind.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ImplementationPlatformKind.kt
index 0f717ce..03252ec 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ImplementationPlatformKind.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ImplementationPlatformKind.kt
@@ -6,7 +6,10 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.util
 
 import org.jetbrains.kotlin.analysis.api.KaImplementationDetail
+import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinProjectStructureProvider
+import org.jetbrains.kotlin.analysis.api.projectStructure.KaModule
 import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.platform.isCommon
 import org.jetbrains.kotlin.platform.isJs
 import org.jetbrains.kotlin.platform.isWasm
 import org.jetbrains.kotlin.platform.jvm.isJvm
@@ -16,7 +19,8 @@
  * Represents the platform kind of implementation (non-common) module.
  * Unlike [TargetPlatform], here we discard specifics of platforms (such as the compatible JDK version).
  */
-internal enum class ImplementationPlatformKind(private val matcher: (TargetPlatform) -> Boolean) {
+@KaImplementationDetail
+enum class ImplementationPlatformKind(private val matcher: (TargetPlatform) -> Boolean) {
     JVM({ it.isJvm() }),
     JAVASCRIPT({ it.isJs() }),
     NATIVE({ it.isNative() }),
@@ -35,4 +39,32 @@
             return entries.firstOrNull { it.matches(targetPlatform) }
         }
     }
+}
+
+/**
+ * Provides an implementation counterpart for common modules.
+ */
+@KaImplementationDetail
+fun interface LLPlatformActualizer {
+    /**
+     * Returns the implementation module for the given [module], or `null` if there is no implementing module, or the given module
+     * is not a common module.
+     */
+    fun actualize(module: KaModule): KaModule?
+}
+
+/**
+ * A simple implementation of [LLPlatformActualizer] which returns the first module of the given [kind].
+ */
+@KaImplementationDetail
+class LLKindBasedPlatformActualizer(private val kind: ImplementationPlatformKind) : LLPlatformActualizer {
+    override fun actualize(module: KaModule): KaModule? {
+        if (!module.targetPlatform.isCommon()) {
+            return null
+        }
+
+        return KotlinProjectStructureProvider.getInstance(module.project)
+            .getImplementingModules(module)
+            .find { ImplementationPlatformKind.fromTargetPlatform(it.targetPlatform) == kind }
+    }
 }
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractCompilationPeerAnalysisTest.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractCompilationPeerAnalysisTest.kt
index a103f66..10a768e 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractCompilationPeerAnalysisTest.kt
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractCompilationPeerAnalysisTest.kt
@@ -9,6 +9,8 @@
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.CompilationPeerCollector
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ImplementationPlatformKind
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLKindBasedPlatformActualizer
 import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedTest
 import org.jetbrains.kotlin.analysis.test.framework.projectStructure.KtTestModule
 import org.jetbrains.kotlin.psi.KtFile
@@ -24,7 +26,9 @@
         val resolveSession = mainModule.ktModule.getFirResolveSession(project)
         val firFile = mainFile.getOrBuildFirFile(resolveSession)
 
-        val compilationPeerData = CompilationPeerCollector.process(firFile)
+        val platformKind = ImplementationPlatformKind.fromTargetPlatform(firFile.moduleData.platform)
+        val actualizer = platformKind?.let(::LLKindBasedPlatformActualizer)
+        val compilationPeerData = CompilationPeerCollector.process(firFile, actualizer)
 
         val filesToCompile = compilationPeerData.peers.values
             .flatten()