FIR IDE: Correctly handle properties in KtFirCompletionCandidateChecker

- Introduce `EnclosingDeclarationContext` to remove code duplication

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt

@@ -20,12 +20,16 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirFunctionSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirPropertySymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.EnclosingDeclarationContext
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.buildCompletionContext
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.fakeEnclosingDeclaration
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.psi.KtCallableDeclaration
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal class KtFirCompletionCandidateChecker(
@@ -34,7 +38,7 @@ internal class KtFirCompletionCandidateChecker(
 ) : KtCompletionCandidateChecker(), KtFirAnalysisSessionComponent {
     override val analysisSession: KtFirAnalysisSession by weakRef(analysisSession)
 
-    private val completionContextCache = HashMap<Pair<FirFile, KtNamedFunction>, LowLevelFirApiFacade.FirCompletionContext>()
+    private val completionContextCache = HashMap<Pair<FirFile, KtCallableDeclaration>, LowLevelFirApiFacade.FirCompletionContext>()
 
     override fun checkExtensionFitsCandidate(
         firSymbolForCandidate: KtCallableSymbol,
@@ -86,17 +90,10 @@ internal class KtFirCompletionCandidateChecker(
         firFile: FirFile,
         fakeNameExpression: KtSimpleNameExpression
     ): Sequence<ImplicitReceiverValue<*>?> {
-        val fakeEnclosingFunction = fakeNameExpression.getNonStrictParentOfType<KtNamedFunction>()
+        val enclosingContext = EnclosingDeclarationContext.detect(originalFile, fakeNameExpression)
-            ?: error("Cannot find enclosing function for ${fakeNameExpression.getElementTextInContext()}")
+
-        val originalEnclosingFunction = originalFile.findFunctionDeclarationAt(fakeEnclosingFunction.textOffset)
+        val completionContext = completionContextCache.getOrCreate(firFile to enclosingContext.fakeEnclosingDeclaration) {
-            ?: error("Cannot find enclosing function for completion in provided position (or position is absent)")
+            enclosingContext.buildCompletionContext(firFile, firResolveState)
-        val completionContext = completionContextCache.getOrCreate(firFile to fakeEnclosingFunction) {
-            LowLevelFirApiFacade.buildCompletionContextForFunction(
-                firFile,
-                fakeEnclosingFunction,
-                originalEnclosingFunction,
-                state = firResolveState
-            )
         }
 
         val towerDataContext = completionContext.getTowerDataContext(fakeNameExpression)
@@ -106,9 +103,4 @@ internal class KtFirCompletionCandidateChecker(
             yieldAll(towerDataContext.implicitReceiverStack)
         }
     }
-
-    private fun KtFile.findFunctionDeclarationAt(offset: Int): KtNamedFunction? =
-        findElementAt(offset)
-            ?.getNonStrictParentOfType<KtNamedFunction>()
-            ?.takeIf { it.textOffset == offset }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirScopeProvider.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.openapi.project.Project
-import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -26,26 +25,18 @@ import org.jetbrains.kotlin.idea.frontend.api.components.KtScopeProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.*
-import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.KtFirDeclaredMemberScope
-import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.KtFirDelegatingScope
-import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.KtFirMemberScope
-import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.KtFirNonStarImportingScope
-import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.KtFirPackageScope
-import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.KtFirStarImportingScope
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.EnclosingDeclarationContext
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.buildCompletionContext
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.scopes.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtProperty
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import java.util.*
 
 internal class KtFirScopeProvider(
@@ -138,41 +129,11 @@ internal class KtFirScopeProvider(
         positionInFakeFile: KtElement
     ): LowLevelFirApiFacade.FirCompletionContext {
         val originalFirFile = originalFile.getOrBuildFirOfType<FirFile>(firResolveState)
-        val fakeEnclosingFunction = positionInFakeFile.getNonStrictParentOfType<KtNamedFunction>()
+        val declarationContext = EnclosingDeclarationContext.detect(originalFile, positionInFakeFile)
 
-        if (fakeEnclosingFunction != null) {
+        return declarationContext.buildCompletionContext(originalFirFile, firResolveState)
-            val originalEnclosingFunction = originalFile.findDeclarationOfTypeAt<KtNamedFunction>(fakeEnclosingFunction.textOffset)
-                ?: error("Cannot find original function matching to ${fakeEnclosingFunction.getElementTextInContext()} in $originalFile")
-
-            return LowLevelFirApiFacade.buildCompletionContextForFunction(
-                originalFirFile,
-                fakeEnclosingFunction,
-                originalEnclosingFunction,
-                state = firResolveState
-            )
-        }
-
-        val fakeEnclosingProperty = positionInFakeFile.parentsOfType<KtProperty>().firstOrNull { !it.isLocal }
-        if (fakeEnclosingProperty != null) {
-            val originalEnclosingProperty = originalFile.findDeclarationOfTypeAt<KtProperty>(fakeEnclosingProperty.textOffset)
-                ?: error("Cannot find original property matching to ${fakeEnclosingProperty.getElementTextInContext()} in $originalFile")
-
-            return LowLevelFirApiFacade.buildCompletionContextForProperty(
-                originalFirFile,
-                fakeEnclosingProperty,
-                originalEnclosingProperty,
-                state = firResolveState
-            )
-        }
-
-        error("Cannot find enclosing declaration for ${positionInFakeFile.getElementTextInContext()}")
     }
 
-    private inline fun<reified T: KtElement> KtFile.findDeclarationOfTypeAt(offset: Int): T? =
-        findElementAt(offset)
-            ?.getNonStrictParentOfType<T>()
-            ?.takeIf { it.textOffset == offset }
-
     private fun convertToKtScope(firScope: FirScope): KtScope {
         firScopeStorage.register(firScope)
         return when (firScope) {

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

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+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.FirModuleResolveState
+import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+
+internal sealed class EnclosingDeclarationContext {
+    companion object {
+        fun detect(originalFile: KtFile, positionInFakeFile: KtElement): EnclosingDeclarationContext {
+            val fakeFunction = positionInFakeFile.getNonStrictParentOfType<KtNamedFunction>()
+            if (fakeFunction != null) {
+                val originalFunction = originalFile.findDeclarationOfTypeAt<KtNamedFunction>(fakeFunction.textOffset)
+                    ?: error("Cannot find original function matching to ${fakeFunction.getElementTextInContext()} in $originalFile")
+
+                return FunctionContext(fakeFunction, originalFunction)
+            }
+
+            val fakeProperty = positionInFakeFile.parentsOfType<KtProperty>().firstOrNull { !it.isLocal }
+            if (fakeProperty != null) {
+                val originalProperty = originalFile.findDeclarationOfTypeAt<KtProperty>(fakeProperty.textOffset)
+                    ?: error("Cannot find original property matching to ${fakeProperty.getElementTextInContext()} in $originalFile")
+
+                return PropertyContext(fakeProperty, originalProperty)
+            }
+
+            error("Cannot find enclosing declaration for ${positionInFakeFile.getElementTextInContext()}")
+        }
+    }
+}
+
+internal class FunctionContext(
+    val fakeEnclosingFunction: KtNamedFunction,
+    val originalEnclosingFunction: KtNamedFunction
+) : EnclosingDeclarationContext()
+
+internal class PropertyContext(
+    val fakeEnclosingProperty: KtProperty,
+    val originalEnclosingProperty: KtProperty
+) : EnclosingDeclarationContext()
+
+internal val EnclosingDeclarationContext.fakeEnclosingDeclaration: KtCallableDeclaration
+    get() = when (this) {
+        is FunctionContext -> fakeEnclosingFunction
+        is PropertyContext -> fakeEnclosingProperty
+    }
+
+internal fun EnclosingDeclarationContext.buildCompletionContext(originalFirFile: FirFile, firResolveState: FirModuleResolveState) =
+    when (this) {
+        is FunctionContext -> LowLevelFirApiFacade.buildCompletionContextForFunction(
+            originalFirFile,
+            fakeEnclosingFunction,
+            originalEnclosingFunction,
+            state = firResolveState
+        )
+
+        is PropertyContext -> LowLevelFirApiFacade.buildCompletionContextForProperty(
+            originalFirFile,
+            fakeEnclosingProperty,
+            originalEnclosingProperty,
+            state = firResolveState
+        )
+    }
+
+private inline fun <reified T : KtElement> KtFile.findDeclarationOfTypeAt(offset: Int): T? =
+    findElementAt(offset)
+        ?.getNonStrictParentOfType<T>()
+        ?.takeIf { it.textOffset == offset }