FIR IDE: simplify completion context retrieving

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

@@ -70,12 +70,10 @@ internal class FirModuleResolveStateForCompletion(
     }
 
     override fun lazyResolveDeclarationForCompletion(
-        firFunction: FirDeclaration,
+        firDeclaration: FirDeclaration,
         containerFirFile: FirFile,
-        firIdeProvider: FirProvider,
-        toPhase: FirResolvePhase,
     ) {
-        originalState.lazyResolveDeclarationForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase)
+        originalState.lazyResolveDeclarationForCompletion(firDeclaration, containerFirFile)
     }
 
     override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? {

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

@@ -157,19 +157,17 @@ internal class FirModuleResolveStateImpl(
     }
 
     override fun lazyResolveDeclarationForCompletion(
-        firFunction: FirDeclaration,
+        firDeclaration: FirDeclaration,
         containerFirFile: FirFile,
-        firIdeProvider: FirProvider,
-        toPhase: FirResolvePhase,
     ) {
         firFileBuilder.runCustomResolveWithPCECheck(containerFirFile, rootModuleSession.cache) {
             firLazyDeclarationResolver.runLazyResolveWithoutLock(
-                firFunction,
+                firDeclaration,
                 rootModuleSession.cache,
                 containerFirFile,
-                firIdeProvider,
-                fromPhase = firFunction.resolvePhase,
-                toPhase,
+                containerFirFile.session.firIdeProvider,
+                fromPhase = firDeclaration.resolvePhase,
+                toPhase = FirResolvePhase.BODY_RESOLVE,
                 towerDataContextCollector = collector,
                 checkPCE = true
             )

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

@@ -0,0 +1,118 @@
+/*
+ * Copyright 2010-2021 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.fir.low.level.api.api
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.builder.RawFirFragmentForLazyBodiesBuilder
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyAccessorCopy
+import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunctionCopy
+import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.KtProperty
+
+object DeclarationCopyBuilder {
+    fun createDeclarationCopy(
+        originalFirDeclaration: FirDeclaration,
+        fakeKtDeclaration: KtDeclaration,
+        state: FirModuleResolveState,
+    ): FirDeclaration {
+        return when (fakeKtDeclaration) {
+            is KtNamedFunction -> buildFunctionCopy(
+                fakeKtDeclaration,
+                originalFirDeclaration as FirSimpleFunction,
+                state
+            )
+            is KtProperty -> buildPropertyCopy(
+                fakeKtDeclaration,
+                originalFirDeclaration as FirProperty,
+                state
+            )
+            else -> error("Unsupported declaration ${fakeKtDeclaration::class.simpleName}")
+        }
+    }
+
+    private fun buildFunctionCopy(
+        element: KtNamedFunction,
+        originalFunction: FirSimpleFunction,
+        state: FirModuleResolveState,
+    ): FirSimpleFunction {
+        val builtFunction = createCopy(element, originalFunction) as FirSimpleFunction
+
+        // right now we can't resolve builtFunction header properly, as it built right in air,
+        // without file, which is now required for running stages other then body resolve, so we
+        // take original function header (which is resolved) and copy replacing body with body from builtFunction
+        return buildSimpleFunctionCopy(originalFunction) {
+            body = builtFunction.body
+            symbol = builtFunction.symbol
+            resolvePhase = minOf(originalFunction.resolvePhase, FirResolvePhase.DECLARATIONS)
+            source = builtFunction.source
+            session = state.rootModuleSession
+        }.apply { reassignAllReturnTargets(builtFunction) }
+    }
+
+    private fun buildPropertyCopy(
+        element: KtProperty,
+        originalProperty: FirProperty,
+        state: FirModuleResolveState
+    ): FirProperty {
+        val builtProperty = createCopy(element, originalProperty) as FirProperty
+
+        val originalSetter = originalProperty.setter
+        val builtSetter = builtProperty.setter
+
+        // setter has a header with `value` parameter, and we want it type to be resolved
+        val copySetter = if (originalSetter != null && builtSetter != null) {
+            buildPropertyAccessorCopy(originalSetter) {
+                body = builtSetter.body
+                symbol = builtSetter.symbol
+                resolvePhase = minOf(builtSetter.resolvePhase, FirResolvePhase.DECLARATIONS)
+                source = builtSetter.source
+                session = state.rootModuleSession
+            }.apply { reassignAllReturnTargets(builtSetter) }
+        } else {
+            builtSetter
+        }
+
+        return org.jetbrains.kotlin.fir.declarations.builder.buildPropertyCopy(originalProperty) {
+            symbol = builtProperty.symbol
+            initializer = builtProperty.initializer
+
+            getter = builtProperty.getter
+            setter = copySetter
+
+            resolvePhase = minOf(originalProperty.resolvePhase, FirResolvePhase.DECLARATIONS)
+            source = builtProperty.source
+            session = state.rootModuleSession
+        }
+    }
+
+    fun createCopy(
+        fakeKtDeclaration: KtDeclaration,
+        originalFirDeclaration: FirDeclaration,
+    ): FirDeclaration {
+        return RawFirFragmentForLazyBodiesBuilder.build(
+            session = originalFirDeclaration.session,
+            baseScopeProvider = originalFirDeclaration.session.firIdeProvider.kotlinScopeProvider,
+            designation = originalFirDeclaration.collectDesignation().fullDesignation,
+            declaration = fakeKtDeclaration
+        )
+    }
+
+    private fun FirFunction<*>.reassignAllReturnTargets(from: FirFunction<*>) {
+        this.accept(object : FirVisitorVoid() {
+            override fun visitElement(element: FirElement) {
+                if (element is FirReturnExpression && element.target.labeledElement == from) {
+                    element.target.bind(this@reassignAllReturnTargets)
+                }
+                element.acceptChildren(this)
+            }
+        })
+    }
+}

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

@@ -83,10 +83,8 @@ abstract class FirModuleResolveState {
 
     // todo temporary, used only in completion
     internal abstract fun lazyResolveDeclarationForCompletion(
-        firFunction: FirDeclaration,
+        firDeclaration: FirDeclaration,
         containerFirFile: FirFile,
-        firIdeProvider: FirProvider,
-        toPhase: FirResolvePhase,
     )
 
     internal abstract fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile?

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

@@ -5,20 +5,11 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
-import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyAccessorCopy
-import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyCopy
-import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunctionCopy
-import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
-import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateForCompletion
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
-import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtProperty
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtFile
 
 object LowLevelFirApiFacadeForCompletion {
     fun getResolveStateForCompletion(originalState: FirModuleResolveState): FirModuleResolveState {
@@ -26,102 +17,17 @@ object LowLevelFirApiFacadeForCompletion {
         return FirModuleResolveStateForCompletion(originalState.project, originalState)
     }
 
-    fun recordCompletionContextForFunction(
+    fun recordCompletionContextForDeclaration(
         firFile: FirFile,
-        fakeElement: KtNamedFunction,
-        originalElement: KtNamedFunction,
+        fakeKtDeclaration: KtDeclaration,
+        originalKtDeclaration: KtDeclaration,
+        fakeContainingFile: KtFile,
         state: FirModuleResolveState,
     ) {
-        val firIdeProvider = firFile.session.firIdeProvider
+        val originalFirDeclaration = originalKtDeclaration.getOrBuildFirOfType<FirDeclaration>(state)
+        val fakeFirDeclaration = DeclarationCopyBuilder.createDeclarationCopy(originalFirDeclaration, fakeKtDeclaration, state)
 
-        val originalFunction = state.getOrBuildFirFor(originalElement) as FirSimpleFunction
-        val copyFunction = buildFunctionCopyForCompletion(firIdeProvider, fakeElement, originalFunction, state)
-
-        state.lazyResolveDeclarationForCompletion(copyFunction, firFile, firIdeProvider, FirResolvePhase.BODY_RESOLVE)
-        state.recordPsiToFirMappingsForCompletionFrom(copyFunction, firFile, fakeElement.containingKtFile)
-    }
-
-    fun recordCompletionContextForProperty(
-        firFile: FirFile,
-        fakeElement: KtProperty,
-        originalElement: KtProperty,
-        state: FirModuleResolveState,
-    ) {
-        val firIdeProvider = firFile.session.firIdeProvider
-
-        val originalProperty = state.getOrBuildFirFor(originalElement) as FirProperty
-        val copyProperty = buildPropertyCopyForCompletion(firIdeProvider, fakeElement, originalProperty, state)
-
-        state.lazyResolveDeclarationForCompletion(copyProperty, firFile, firIdeProvider, FirResolvePhase.BODY_RESOLVE)
-        state.recordPsiToFirMappingsForCompletionFrom(copyProperty, firFile, fakeElement.containingKtFile)
-    }
-
-    private fun buildFunctionCopyForCompletion(
-        firIdeProvider: FirIdeProvider,
-        element: KtNamedFunction,
-        originalFunction: FirSimpleFunction,
-        state: FirModuleResolveState,
-    ): FirSimpleFunction {
-        val builtFunction = firIdeProvider.buildFunctionWithBody(element, originalFunction)
-
-        // right now we can't resolve builtFunction header properly, as it built right in air,
-        // without file, which is now required for running stages other then body resolve, so we
-        // take original function header (which is resolved) and copy replacing body with body from builtFunction
-        return buildSimpleFunctionCopy(originalFunction) {
-            body = builtFunction.body
-            symbol = builtFunction.symbol as FirNamedFunctionSymbol
-            resolvePhase = minOf(originalFunction.resolvePhase, FirResolvePhase.DECLARATIONS)
-            source = builtFunction.source
-            session = state.rootModuleSession
-        }.apply { reassignAllReturnTargets(builtFunction) }
-    }
-
-    private fun buildPropertyCopyForCompletion(
-        firIdeProvider: FirIdeProvider,
-        element: KtProperty,
-        originalProperty: FirProperty,
-        state: FirModuleResolveState
-    ): FirProperty {
-        val builtProperty = firIdeProvider.buildPropertyWithBody(element, originalProperty)
-
-        val originalSetter = originalProperty.setter
-        val builtSetter = builtProperty.setter
-
-        // setter has a header with `value` parameter, and we want it type to be resolved
-        val copySetter = if (originalSetter != null && builtSetter != null) {
-            buildPropertyAccessorCopy(originalSetter) {
-                body = builtSetter.body
-                symbol = builtSetter.symbol
-                resolvePhase = minOf(builtSetter.resolvePhase, FirResolvePhase.DECLARATIONS)
-                source = builtSetter.source
-                session = state.rootModuleSession
-            }.apply { reassignAllReturnTargets(builtSetter) }
-        } else {
-            builtSetter
-        }
-
-        return buildPropertyCopy(originalProperty) {
-            symbol = builtProperty.symbol
-            initializer = builtProperty.initializer
-
-            getter = builtProperty.getter
-            setter = copySetter
-
-            resolvePhase = minOf(originalProperty.resolvePhase, FirResolvePhase.DECLARATIONS)
-            source = builtProperty.source
-            session = state.rootModuleSession
-        }
-    }
-
-
-    private fun FirFunction<*>.reassignAllReturnTargets(from: FirFunction<*>) {
-        this.accept(object : FirVisitorVoid() {
-            override fun visitElement(element: FirElement) {
-                if (element is FirReturnExpression && element.target.labeledElement == from) {
-                    element.target.bind(this@reassignAllReturnTargets)
-                }
-                element.acceptChildren(this)
-            }
-        })
+        state.lazyResolveDeclarationForCompletion(fakeFirDeclaration, firFile)
+        state.recordPsiToFirMappingsForCompletionFrom(fakeFirDeclaration, firFile, fakeContainingFile)
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.DeclarationCopyBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileDiagnosticRetriever
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileStructureElementDiagnostics
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
@@ -79,7 +80,7 @@ internal class ReanalyzableFunctionStructureElement(
         towerDataContextCollector: FirTowerDataContextCollector,
     ): ReanalyzableFunctionStructureElement {
         val originalFunction = firSymbol.fir as FirSimpleFunction
-        val newFunction = firIdeProvider.buildFunctionWithBody(newKtDeclaration, originalFunction) as FirSimpleFunction
+        val newFunction = DeclarationCopyBuilder.createCopy(newKtDeclaration, originalFunction) as FirSimpleFunction
 
         return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalFunction, newFunction) {
             firLazyDeclarationResolver.lazyResolveDeclaration(
@@ -121,7 +122,7 @@ internal class ReanalyzablePropertyStructureElement(
         towerDataContextCollector: FirTowerDataContextCollector,
     ): ReanalyzablePropertyStructureElement {
         val originalProperty = firSymbol.fir
-        val newProperty = firIdeProvider.buildPropertyWithBody(newKtDeclaration, originalProperty)
+        val newProperty = DeclarationCopyBuilder.createCopy(newKtDeclaration, originalProperty) as FirProperty
 
         return FileStructureUtil.withDeclarationReplaced(firFile, cache, originalProperty, newProperty) {
             firLazyDeclarationResolver.lazyResolveDeclaration(

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

@@ -91,32 +91,6 @@ internal class FirIdeProvider(
     override fun getFirFilesByPackage(fqName: FqName): List<FirFile> = error("Should not be called in FIR IDE")
 
 
-    // TODO move out of here
-    // used only for completion
-    fun buildFunctionWithBody(
-        ktNamedFunction: KtNamedFunction,
-        original: FirFunction<*>,
-    ): FirFunction<*> {
-        return RawFirFragmentForLazyBodiesBuilder.build(
-            session = original.session,
-            baseScopeProvider = original.session.firIdeProvider.kotlinScopeProvider,
-            designation = original.collectDesignation().fullDesignation,
-            declaration = ktNamedFunction
-        ) as FirFunction<*>
-    }
-
-    fun buildPropertyWithBody(
-        ktNamedFunction: KtProperty,
-        original: FirProperty,
-    ): FirProperty {
-        return RawFirFragmentForLazyBodiesBuilder.build(
-            session = original.session,
-            baseScopeProvider = original.session.firIdeProvider.kotlinScopeProvider,
-            designation = original.collectDesignation().fullDesignation,
-            declaration = ktNamedFunction
-        ) as FirProperty
-    }
-
     @FirProviderInternals
     override fun recordGeneratedClass(owner: FirAnnotatedDeclaration, klass: FirRegularClass) {
         TODO()

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.utils
 
+import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
@@ -13,30 +14,30 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import kotlin.reflect.KClass
 
 internal sealed class EnclosingDeclarationContext {
+    abstract val fakeKtEnclosingDeclaration: KtDeclaration
+    abstract val originalKtEnclosingDeclaration: KtDeclaration
+
     companion object {
+        private fun KtDeclaration.canBeEnclosingDeclaration(): Boolean = when (this) {
+            is KtNamedFunction -> isTopLevel || containingClassOrObject?.isLocal == false
+            is KtProperty -> isTopLevel || containingClassOrObject?.isLocal == false
+            is KtClassOrObject -> !isLocal
+            is KtTypeAlias -> isTopLevel() || containingClassOrObject?.isLocal == false
+            else -> false
+        }
+
         fun detect(originalFile: KtFile, positionInFakeFile: KtElement): EnclosingDeclarationContext {
-            val fakeFunction = positionInFakeFile.parentsOfType<KtNamedFunction>().firstOrNull {
-                it.isTopLevel || it.containingClassOrObject?.isLocal == false
+            val fakeKtDeclaration = positionInFakeFile.parentsOfType<KtNamedDeclaration>().firstOrNull { ktDeclaration ->
+                ktDeclaration.canBeEnclosingDeclaration()
             }
-            if (fakeFunction != null) {
-                val originalFunction = originalFile.findDeclarationOfTypeAt<KtNamedFunction>(fakeFunction.textOffset)
-                    ?: error("Cannot find original function matching to ${fakeFunction.getElementTextInContext()} in $originalFile")
-                recordOriginalDeclaration(originalFunction, fakeFunction)
-                return FunctionContext(fakeFunction, originalFunction)
-            }
-
-            val fakeProperty = positionInFakeFile.parentsOfType<KtProperty>().firstOrNull {
-                it.isTopLevel || it.containingClassOrObject?.isLocal == false
-            }
-
-            if (fakeProperty != null) {
-                val originalProperty = originalFile.findDeclarationOfTypeAt<KtProperty>(fakeProperty.textOffset)
-                    ?: error("Cannot find original property matching to ${fakeProperty.getElementTextInContext()} in $originalFile")
-                recordOriginalDeclaration(originalProperty, fakeProperty)
-                return PropertyContext(fakeProperty, originalProperty)
+            if (fakeKtDeclaration != null) {
+                val originalDeclaration = originalFile.findDeclarationOfTypeAt(fakeKtDeclaration.textOffset, fakeKtDeclaration::class)
+                    ?: error("Cannot find original function matching to ${fakeKtDeclaration.getElementTextInContext()} in $originalFile")
+                recordOriginalDeclaration(originalDeclaration, fakeKtDeclaration)
+                return EnclosingDeclarationContextImpl(fakeKtDeclaration, originalDeclaration)
             }
 
             error("Cannot find enclosing declaration for ${positionInFakeFile.getElementTextInContext()}")
@@ -55,41 +56,24 @@ internal sealed class EnclosingDeclarationContext {
     }
 }
 
-internal class FunctionContext(
-    val fakeEnclosingFunction: KtNamedFunction,
-    val originalEnclosingFunction: KtNamedFunction
+internal class EnclosingDeclarationContextImpl(
+    override val fakeKtEnclosingDeclaration: KtDeclaration,
+    override val originalKtEnclosingDeclaration: KtDeclaration
 ) : 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.recordCompletionContext(originalFirFile: FirFile, firResolveState: FirModuleResolveState) {
-    when (this) {
-        is FunctionContext -> LowLevelFirApiFacadeForCompletion.recordCompletionContextForFunction(
-            originalFirFile,
-            fakeEnclosingFunction,
-            originalEnclosingFunction,
-            state = firResolveState
-        )
-
-        is PropertyContext -> LowLevelFirApiFacadeForCompletion.recordCompletionContextForProperty(
-            originalFirFile,
-            fakeEnclosingProperty,
-            originalEnclosingProperty,
-            state = firResolveState
-        )
-    }
+    LowLevelFirApiFacadeForCompletion.recordCompletionContextForDeclaration(
+        originalFirFile,
+        fakeKtEnclosingDeclaration,
+        originalKtEnclosingDeclaration,
+        state = firResolveState,
+        fakeContainingFile = fakeKtEnclosingDeclaration.containingKtFile
+    )
 }
 
-private inline fun <reified T : KtElement> KtFile.findDeclarationOfTypeAt(offset: Int): T? =
-    findElementAt(offset)
-        ?.getNonStrictParentOfType<T>()
-        ?.takeIf { it.textOffset == offset }
+private fun <T : KtElement> KtFile.findDeclarationOfTypeAt(offset: Int, declarationType: KClass<T>): T? {
+    val elementAtOffset = findElementAt(offset) ?: return null
+    return PsiTreeUtil.getParentOfType(elementAtOffset, declarationType.java, false)?.takeIf { it.textOffset == offset }
+}