FIR IDE: fix working with copied file in completion

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt

@@ -9,6 +9,7 @@ import com.intellij.codeInsight.completion.*
 import com.intellij.patterns.PlatformPatterns
 import com.intellij.patterns.PsiJavaPatterns
 import com.intellij.util.ProcessingContext
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
 import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.getAnalysisSessionFor
@@ -71,9 +72,16 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
         lookupElementFactory.createLookupElement(symbol)?.let(::addElement)
     }
 
+    private fun recordOriginalFile(completionParameters: CompletionParameters) {
+        val originalFile = completionParameters.originalFile as? KtFile ?: return
+        val fakeFile = completionParameters.position.containingFile as? KtFile ?: return
+        fakeFile.originalKtFile = originalFile
+    }
+
     @OptIn(InvalidWayOfUsingAnalysisSession::class)
     fun addCompletions(parameters: CompletionParameters, result: CompletionResultSet) {
         val originalFile = parameters.originalFile as? KtFile ?: return
+        recordOriginalFile(parameters)
 
         val reference = (parameters.position.parent as? KtSimpleNameExpression)?.mainReference ?: return
         val nameExpression = reference.expression.takeIf { it !is KtLabelReferenceExpression } ?: return

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt

@@ -12,12 +12,15 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.containingKtFileIfAny
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -74,9 +77,13 @@ internal class FirModuleResolveStateForCompletion(
     }
 
     override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? {
-        return cache.getContainerFirFile(declaration)
+        val ktFile = declaration.containingKtFileIfAny ?: return null
+        cache.getCachedFirFile(ktFile)?.let { return it }
+        ktFile.originalKtFile?.let(cache::getCachedFirFile)?.let { return it }
+        return null
     }
 
+
     override fun getDiagnostics(element: KtElement): List<Diagnostic> {
         error("Diagnostics should not be retrieved in completion")
     }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/declarationUtils.kt

@@ -42,8 +42,9 @@ internal fun KtDeclaration.findFirDeclarationForAnyFirSourceDeclaration(
     val nonLocalDeclaration = getNonLocalContainingOrThisDeclaration()
         ?.findSourceNonLocalFirDeclaration(firFileBuilder, firSymbolProvider, moduleFileCache)
         ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile, moduleFileCache, lazyBodiesMode = true)
+    val originalDeclaration = originalDeclaration
     val fir = FirElementFinder.findElementIn<FirDeclaration>(nonLocalDeclaration) { firDeclaration ->
-        firDeclaration.psi == this
+        firDeclaration.psi == this || firDeclaration.psi == originalDeclaration
     }
     return fir
         ?: error("FirDeclaration was not found for\n${getElementTextInContext()}")
@@ -105,9 +106,11 @@ private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider(
 }
 
 val ORIGINAL_DECLARATION_KEY = com.intellij.openapi.util.Key<KtDeclaration>("ORIGINAL_DECLARATION_KEY")
-
 var KtDeclaration.originalDeclaration by UserDataProperty(ORIGINAL_DECLARATION_KEY)
 
+private val ORIGINAL_KT_FILE_KEY = com.intellij.openapi.util.Key<KtFile>("ORIGINAL_KT_FILE_KEY")
+var KtFile.originalKtFile by UserDataProperty(ORIGINAL_KT_FILE_KEY)
+
 
 private fun KtClassOrObject.findFir(firSymbolProvider: FirSymbolProvider): FirRegularClass? {
     val classId = classIdIfNonLocal() ?: return null

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/utils.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtFile
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.locks.Lock
 
@@ -63,6 +64,9 @@ internal val FirDeclaration.ktDeclaration: KtDeclaration
         return psi as KtDeclaration
     }
 
+internal val FirDeclaration.containingKtFileIfAny: KtFile?
+    get() = psi?.containingFile as? KtFile
+
 
 internal fun IdeaModuleInfo.collectTransitiveDependenciesWithSelf(): List<IdeaModuleInfo> {
     val result = mutableSetOf<IdeaModuleInfo>()

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/EnclosingDeclarationContext.kt

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.utils
 import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForCompletion
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
@@ -22,7 +21,7 @@ internal sealed class EnclosingDeclarationContext {
             if (fakeFunction != null) {
                 val originalFunction = originalFile.findDeclarationOfTypeAt<KtNamedFunction>(fakeFunction.textOffset)
                     ?: error("Cannot find original function matching to ${fakeFunction.getElementTextInContext()} in $originalFile")
-                fakeFunction.originalDeclaration = originalFunction
+                recordOriginalDeclaration(originalFunction, fakeFunction)
                 return FunctionContext(fakeFunction, originalFunction)
             }
 
@@ -30,13 +29,23 @@ internal sealed class EnclosingDeclarationContext {
             if (fakeProperty != null) {
                 val originalProperty = originalFile.findDeclarationOfTypeAt<KtProperty>(fakeProperty.textOffset)
                     ?: error("Cannot find original property matching to ${fakeProperty.getElementTextInContext()} in $originalFile")
-                fakeProperty.originalDeclaration = originalProperty
-
+                recordOriginalDeclaration(originalProperty, fakeProperty)
                 return PropertyContext(fakeProperty, originalProperty)
             }
 
             error("Cannot find enclosing declaration for ${positionInFakeFile.getElementTextInContext()}")
         }
+
+        private fun recordOriginalDeclaration(originalDeclaration: KtNamedDeclaration, fakeDeclaration: KtNamedDeclaration) {
+            require(!fakeDeclaration.isPhysical)
+            require(originalDeclaration.containingKtFile !== fakeDeclaration.containingKtFile)
+            val originalDeclrationParents = originalDeclaration.parentsOfType<KtDeclaration>().toList()
+            val fakeDeclarationParents = fakeDeclaration.parentsOfType<KtDeclaration>().toList()
+
+            originalDeclrationParents.zip(fakeDeclarationParents) { original, fake ->
+                fake.originalDeclaration = original
+            }
+        }
     }
 }