From 70b9af05665a45f07a79512e2829705d684851f9 Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Wed, 7 Jun 2023 09:50:16 +0200
Subject: [PATCH] [LL] don't request provider with non-physical element

Otherwise, if fir was requested with non-physical element
 inside existing analysis session (not on-air, e.g.,
`FirPositionCompletionContextDetector#analyzeInContext` with
 `FirSimpleParameterPositionContext`) all consequent requests
 would return fir with that non-physical source,
 which might be already invalidated.

 On-air sessions should be protected from invalidation problem but some clients
 might expect original elements anyway and in this case, it would be
 impossible to retrieve correct fir for them.

 In order to find parameters, allow to traverse tree till functions.
---
 .../api/fir/providers/LLFirProviderHelper.kt  |  2 ++
 .../level/api/fir/util/declarationUtils.kt    | 19 ++++++++++++++-----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt
index 2a3b629d1d1..0d556e50a4d 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt
@@ -72,6 +72,7 @@ internal class LLFirProviderHelper(
 
     private val classifierByClassId =
         firSession.firCachesFactory.createCache<ClassId, FirClassLikeDeclaration?, KtClassLikeDeclaration?> { classId, context ->
+            require(context == null || context.isPhysical)
             val ktClass = context ?: declarationProvider.getClassLikeDeclarationByClassId(classId) ?: return@createCache null
 
             if (ktClass.getClassId() == null) return@createCache null
@@ -82,6 +83,7 @@ internal class LLFirProviderHelper(
 
     private val callablesByCallableId =
         firSession.firCachesFactory.createCache<CallableId, List<FirCallableSymbol<*>>, Collection<KtFile>?> { callableId, context ->
+            require(context == null || context.all { it.isPhysical })
             val files = context ?: declarationProvider.getTopLevelCallableFiles(callableId).ifEmpty { return@createCache emptyList() }
             buildList {
                 files.forEach { ktFile ->
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 1b19f780b43..89ef3e7443a 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
@@ -32,7 +32,12 @@ internal fun KtDeclaration.findSourceNonLocalFirDeclaration(
     //TODO test what way faster
     findSourceNonLocalFirDeclarationByProvider(firFileBuilder, provider, containerFirFile)?.let { return it }
     findSourceByTraversingWholeTree(firFileBuilder, containerFirFile)?.let { return it }
-    errorWithFirSpecificEntries("No fir element was found for", psi = this)
+    errorWithFirSpecificEntries(
+        "No fir element was found for",
+        psi = this,
+        fir = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile),
+        additionalInfos = { withEntry("isPhysical", isPhysical.toString()) }
+    )
 }
 
 internal fun KtDeclaration.findFirDeclarationForAnyFirSourceDeclaration(
@@ -68,7 +73,7 @@ internal fun KtElement.findSourceByTraversingWholeTree(
     val isDeclaration = this is KtDeclaration
     return FirElementFinder.findElementIn(
         firFile,
-        canGoInside = { it is FirRegularClass || it is FirScript },
+        canGoInside = { it is FirRegularClass || it is FirScript || it is FirFunction },
         predicate = { firDeclaration ->
             firDeclaration.psi == this || isDeclaration && firDeclaration.psi == originalDeclaration
         }
@@ -80,6 +85,11 @@ private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider(
     provider: FirProvider,
     containerFirFile: FirFile?,
 ): FirDeclaration? {
+    if (!isPhysical) {
+        //do not request providers with non-physical psi in order not to leak them there and 
+        //to avoid inconsistency between physical psi and its copy during completion
+        return null
+    }
     val candidate = when {
         this is KtClassOrObject -> findFir(provider)
         this is KtNamedDeclaration && (this is KtProperty || this is KtNamedFunction) -> {
@@ -97,19 +107,18 @@ private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider(
                     firFile.declarations
                 }
             }
-            val original = originalDeclaration
 
             /*
             It is possible that we will not be able to find needed declaration here when the code is invalid,
             e.g, we have two conflicting declarations with the same name and we are searching in the wrong one
              */
-            declarations?.firstOrNull { it.psi == this || it.psi == original }
+            declarations?.firstOrNull { it.psi == this }
         }
         this is KtConstructor<*> || this is KtClassInitializer -> {
             val containingClass = containingClassOrObject
                 ?: errorWithFirSpecificEntries("Container class should be not null for KtConstructor", psi = this)
             val containerClassFir = containingClass.findFir(provider) as? FirRegularClass ?: return null
-            containerClassFir.declarations.firstOrNull { it.psi === this }
+            containerClassFir.declarations.firstOrNull { it.psi == this }
         }
         this is KtTypeAlias -> findFir(provider)
         this is KtDestructuringDeclaration -> {