FIR IDE: Add building completion context for properties

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirPropertyBuilder.kt

@@ -90,3 +90,31 @@ inline fun buildProperty(init: FirPropertyBuilder.() -> Unit): FirProperty {
     }
     return FirPropertyBuilder().apply(init).build()
 }
+
+@OptIn(ExperimentalContracts::class)
+inline fun buildPropertyCopy(original: FirProperty, init: FirPropertyBuilder.() -> Unit): FirProperty {
+    contract {
+        callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
+    }
+    val copyBuilder = FirPropertyBuilder()
+    copyBuilder.source = original.source
+    copyBuilder.session = original.session
+    copyBuilder.resolvePhase = original.resolvePhase
+    copyBuilder.origin = original.origin
+    copyBuilder.returnTypeRef = original.returnTypeRef
+    copyBuilder.receiverTypeRef = original.receiverTypeRef
+    copyBuilder.name = original.name
+    copyBuilder.initializer = original.initializer
+    copyBuilder.delegate = original.delegate
+    copyBuilder.delegateFieldSymbol = original.delegateFieldSymbol
+    copyBuilder.isVar = original.isVar
+    copyBuilder.getter = original.getter
+    copyBuilder.setter = original.setter
+    copyBuilder.annotations.addAll(original.annotations)
+    copyBuilder.typeParameters.addAll(original.typeParameters)
+    copyBuilder.containerSource = original.containerSource
+    copyBuilder.symbol = original.symbol
+    copyBuilder.isLocal = original.isLocal
+    copyBuilder.status = original.status
+    return copyBuilder.apply(init).build()
+}

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/BuilderConfigurator.kt

@@ -172,6 +172,7 @@ object BuilderConfigurator : AbstractBuilderConfigurator<FirTreeBuilder>(FirTree
             parents += typeParametersOwnerBuilder
             defaultNull("getter", "setter", "containerSource", "delegateFieldSymbol")
             default("resolvePhase", "FirResolvePhase.RAW_FIR")
+            withCopy()
         }
 
         builder(typeOperatorCall) {

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

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.fir.FirElement
 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.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
@@ -46,8 +45,8 @@ abstract class FirModuleResolveState {
     abstract fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D
 
     // todo temporary, used only in completion
-    internal abstract fun lazyResolveFunctionForCompletion(
+    internal abstract fun lazyResolveDeclarationForCompletion(
-        firFunction: FirFunction<*>,
+        firFunction: FirDeclaration,
         containerFirFile: FirFile,
         firIdeProvider: FirIdeProvider,
         toPhase: FirResolvePhase,
@@ -91,8 +90,8 @@ internal open class FirModuleResolveStateImpl(
     override fun getCachedMappingForCompletion(element: KtElement): FirElement? =
         psiToFirCache.getCachedMapping(element)
 
-    override fun lazyResolveFunctionForCompletion(
+    override fun lazyResolveDeclarationForCompletion(
-        firFunction: FirFunction<*>,
+        firFunction: FirDeclaration,
         containerFirFile: FirFile,
         firIdeProvider: FirIdeProvider,
         toPhase: FirResolvePhase,

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

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.fir.FirElement
 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.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
@@ -55,14 +54,14 @@ internal class FirModuleResolveStateForCompletion(
         return originalState.resolvedFirToPhase(declaration, toPhase)
     }
 
-    override fun lazyResolveFunctionForCompletion(
+    override fun lazyResolveDeclarationForCompletion(
-        firFunction: FirFunction<*>,
+        firFunction: FirDeclaration,
         containerFirFile: FirFile,
         firIdeProvider: FirIdeProvider,
         toPhase: FirResolvePhase,
         towerDataContextCollector: FirTowerDataContextCollector
     ) {
-        originalState.lazyResolveFunctionForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextCollector)
+        originalState.lazyResolveDeclarationForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextCollector)
     }
 
     override fun getDiagnostics(element: KtElement): List<Diagnostic> {

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

@@ -10,9 +10,11 @@ import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunctionCopy
+import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyCopy
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
@@ -22,6 +24,7 @@ import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.psi.KtProperty
 
 object LowLevelFirApiFacade {
     fun getResolveStateFor(element: KtElement): FirModuleResolveState =
@@ -82,7 +85,40 @@ object LowLevelFirApiFacade {
         }
 
         val function = frankensteinFunction.apply {
-            state.lazyResolveFunctionForCompletion(this, firFile, firIdeProvider, phase, contextCollector)
+            state.lazyResolveDeclarationForCompletion(this, firFile, firIdeProvider, phase, contextCollector)
+            state.recordPsiToFirMappingsForCompletionFrom(this, firFile, element.containingKtFile)
+        }
+
+        return FirCompletionContext(
+            function.session,
+            contextCollector,
+            state
+        )
+    }
+
+    fun buildCompletionContextForProperty(
+        firFile: FirFile,
+        element: KtProperty,
+        originalElement: KtProperty,
+        state: FirModuleResolveState,
+        phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
+    ): FirCompletionContext {
+        val firIdeProvider = firFile.session.firIdeProvider
+        val originalProperty = state.getOrBuildFirFor(originalElement, phase) as FirProperty
+        val builtProperty = firIdeProvider.buildPropertyWithBody(element)
+        val contextCollector = FirTowerDataContextCollector()
+
+        val frankensteinProperty = buildPropertyCopy(originalProperty) {
+            symbol = builtProperty.symbol
+            getter = builtProperty.getter
+            setter = builtProperty.setter
+            resolvePhase = minOf(originalProperty.resolvePhase, FirResolvePhase.DECLARATIONS)
+            source = builtProperty.source
+            session = state.firIdeSourcesSession
+        }
+
+        val function = frankensteinProperty.apply {
+            state.lazyResolveDeclarationForCompletion(this, firFile, firIdeProvider, phase, contextCollector)
             state.recordPsiToFirMappingsForCompletionFrom(this, firFile, element.containingKtFile)
         }
 

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

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.KtProperty
 
 @ThreadSafeMutableState
 internal class FirIdeProvider(
@@ -103,6 +104,10 @@ internal class FirIdeProvider(
         return RawFirBuilder(session, kotlinScopeProvider, stubMode = false).buildFunctionWithBody(ktNamedFunction)
     }
 
+    fun buildPropertyWithBody(ktNamedFunction: KtProperty): FirProperty {
+        return RawFirBuilder(session, kotlinScopeProvider, stubMode = false).buildPropertyWithBody(ktNamedFunction)
+    }
+
 
     @FirProviderInternals
     override fun recordGeneratedClass(owner: FirAnnotatedDeclaration, klass: FirRegularClass) {

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

@@ -6,7 +6,7 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.openapi.project.Project
-import com.intellij.psi.PsiElement
+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
@@ -44,6 +44,7 @@ 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.*
 
@@ -108,18 +109,7 @@ internal class KtFirScopeProvider(
         originalFile: KtFile,
         positionInFakeFile: KtElement
     ): KtScopeContext = withValidityAssertion {
-        val originalFirFile = originalFile.getOrBuildFirOfType<FirFile>(firResolveState)
+        val completionContext = buildCompletionContextForEnclosingDeclaration(originalFile, positionInFakeFile)
-        val fakeEnclosingFunction = positionInFakeFile.getNonStrictParentOfType<KtNamedFunction>()
-            ?: error("Cannot find enclosing function for ${positionInFakeFile.getElementTextInContext()}")
-        val originalEnclosingFunction = originalFile.findFunctionDeclarationAt(fakeEnclosingFunction.textOffset)
-            ?: error("Cannot find original function matching to ${fakeEnclosingFunction.getElementTextInContext()} in $originalFile")
-
-        val completionContext = LowLevelFirApiFacade.buildCompletionContextForFunction(
-            originalFirFile,
-            fakeEnclosingFunction,
-            originalEnclosingFunction,
-            state = firResolveState
-        )
 
         val towerDataContext = completionContext.getTowerDataContext(positionInFakeFile)
 
@@ -143,9 +133,44 @@ internal class KtFirScopeProvider(
         )
     }
 
-    private fun KtFile.findFunctionDeclarationAt(offset: Int): KtNamedFunction? =
+    private fun buildCompletionContextForEnclosingDeclaration(
+        originalFile: KtFile,
+        positionInFakeFile: KtElement
+    ): LowLevelFirApiFacade.FirCompletionContext {
+        val originalFirFile = originalFile.getOrBuildFirOfType<FirFile>(firResolveState)
+        val fakeEnclosingFunction = positionInFakeFile.getNonStrictParentOfType<KtNamedFunction>()
+
+        if (fakeEnclosingFunction != null) {
+            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<KtNamedFunction>()
+            ?.getNonStrictParentOfType<T>()
             ?.takeIf { it.textOffset == offset }
 
     private fun convertToKtScope(firScope: FirScope): KtScope {