From 297920fde949815ec23ee9f0025da455a71e5a5b Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Wed, 22 Nov 2023 20:31:26 +0900
Subject: [PATCH] [Analysis API] Support dangling file copies in
 'KtFirVisibilityChecker'

Dangling files are copies of files in the context module, or just
depend on it. In any way, dangling file modules see all internals of
the context module.
---
 .../fir/components/KtFirVisibilityChecker.kt  | 33 ++++++++++++-------
 .../level/api/fir/util/declarationUtils.kt    | 10 +++---
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirVisibilityChecker.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirVisibilityChecker.kt
index e79054a9857..22769002e4e 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirVisibilityChecker.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirVisibilityChecker.kt
@@ -16,7 +16,9 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtFileSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithVisibility
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirSafe
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.collectContainingDeclarationsIfNonLocal
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.collectUseSiteContainers
+import org.jetbrains.kotlin.analysis.project.structure.KtDanglingFileModule
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.descriptors.java.JavaVisibilities
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
@@ -48,22 +50,31 @@ internal class KtFirVisibilityChecker(
             candidateSymbol.isVisibleByPsi(useSiteFile)?.let { return it }
         }
 
-        val useSiteFirFile = useSiteFile.firSymbol.fir
-        val containers = position.collectContainingDeclarationsIfNonLocal(analysisSession.firResolveSession).orEmpty()
-
         val dispatchReceiverCanBeExplicit = candidateSymbol is KtCallableSymbol && !candidateSymbol.isExtension
         val explicitDispatchReceiver = runIf(dispatchReceiverCanBeExplicit) {
             receiverExpression?.getOrBuildFirSafe<FirExpression>(analysisSession.firResolveSession)
         }
 
-        val candidateFirSymbol = candidateSymbol.firSymbol.fir as FirMemberDeclaration
+        val candidateDeclaration = candidateSymbol.firSymbol.fir as FirMemberDeclaration
 
-        val rootModuleSession = rootModuleSession
-        return rootModuleSession.visibilityChecker.isVisible(
-            candidateFirSymbol,
-            rootModuleSession,
-            useSiteFirFile,
-            containers,
+        val positionModule = firResolveSession.moduleProvider.getModule(position)
+        val candidateModule = candidateDeclaration.llFirModuleData.ktModule
+
+        val effectiveSession = if (positionModule is KtDanglingFileModule && candidateModule != positionModule) {
+            @Suppress("USELESS_CAST") // Smart cast is only available in K2
+            val contextModule = (positionModule as KtDanglingFileModule).contextModule
+            firResolveSession.getSessionFor(contextModule)
+        } else {
+            firResolveSession.getSessionFor(positionModule)
+        }
+
+        val effectiveContainers = collectUseSiteContainers(position, firResolveSession).orEmpty()
+
+        return effectiveSession.visibilityChecker.isVisible(
+            candidateDeclaration,
+            effectiveSession,
+            useSiteFile.firSymbol.fir,
+            effectiveContainers,
             explicitDispatchReceiver
         )
     }
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt
index 225f94fcee4..0c6f39bf8d9 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt
@@ -90,11 +90,11 @@ internal fun KtDeclaration.findSourceNonLocalFirDeclaration(firFile: FirFile, pr
 }
 
 @KtAnalysisApiInternals
-fun PsiElement.collectContainingDeclarationsIfNonLocal(session: LLFirResolveSession): List<FirDeclaration>? {
-    val ktFile = containingFile as? KtFile ?: return null
-    val ktDeclaration = getNonLocalContainingOrThisDeclaration { !it.canBePartOfParentDeclaration } ?: return null
-    val firFile = session.getOrBuildFirFile(ktFile)
-    return FirElementFinder.findPathToDeclarationWithTarget(firFile, ktDeclaration)
+fun collectUseSiteContainers(element: PsiElement, resolveSession: LLFirResolveSession): List<FirDeclaration>? {
+    val containingDeclaration = element.getNonLocalContainingOrThisDeclaration { !it.canBePartOfParentDeclaration } ?: return null
+    val containingFile = containingDeclaration.containingKtFile
+    val firFile = resolveSession.getOrBuildFirFile(containingFile)
+    return FirElementFinder.findPathToDeclarationWithTarget(firFile, containingDeclaration)
 }
 
 internal fun KtElement.findSourceByTraversingWholeTree(