FIR IDE: Fix completion in case of function with parameters

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

@@ -85,3 +85,27 @@ inline fun buildSimpleFunction(init: FirSimpleFunctionBuilder.() -> Unit): FirSi
     }
     return FirSimpleFunctionBuilder().apply(init).build()
 }
+
+@OptIn(ExperimentalContracts::class)
+inline fun buildSimpleFunctionCopy(original: FirSimpleFunction, init: FirSimpleFunctionBuilder.() -> Unit): FirSimpleFunction {
+    contract {
+        callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
+    }
+    val copyBuilder = FirSimpleFunctionBuilder()
+    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.valueParameters.addAll(original.valueParameters)
+    copyBuilder.body = original.body
+    copyBuilder.status = original.status
+    copyBuilder.containerSource = original.containerSource
+    copyBuilder.contractDescription = original.contractDescription
+    copyBuilder.name = original.name
+    copyBuilder.symbol = original.symbol
+    copyBuilder.annotations.addAll(original.annotations)
+    copyBuilder.typeParameters.addAll(original.typeParameters)
+    return copyBuilder.apply(init).build()
+}

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

@@ -269,6 +269,7 @@ object BuilderConfigurator : AbstractBuilderConfigurator<FirTreeBuilder>(FirTree
             default("contractDescription", "FirEmptyContractDescription")
             useTypes(emptyContractDescriptionType)
             openBuilder()
+            withCopy()
         }
 
         builder(tryExpression) {

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

@@ -43,7 +43,7 @@ private object KotlinHighLevelApiContributor : CompletionProvider<CompletionPara
         val possibleReceiver = nameExpression.getQualifiedExpressionForSelector()?.receiverExpression
 
         with(getAnalysisSessionFor(originalFile).createContextDependentCopy()) {
-            val (implicitScopes, implicitReceivers) = originalFile.getScopeContextForPosition(nameExpression)
+            val (implicitScopes, implicitReceivers) = originalFile.getScopeContextForPosition(parameters.originalPosition, nameExpression)
 
             val typeOfPossibleReceiver = possibleReceiver?.getKtType()
             val possibleReceiverScope = typeOfPossibleReceiver?.let { it.getTypeScope() }

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSession.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.frontend.api
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.frontend.api.components.*
 import org.jetbrains.kotlin.idea.frontend.api.scopes.*
@@ -68,7 +69,8 @@ abstract class KtAnalysisSession(override val token: ValidityToken) : ValidityTo
 
     fun KtType.getTypeScope(): KtScope? = scopeProvider.getTypeScope(this)
 
-    fun KtFile.getScopeContextForPosition(positionInFakeFile: KtElement): KtScopeContext = scopeProvider.getScopeContextForPosition(this, positionInFakeFile)
+    fun KtFile.getScopeContextForPosition(originalPosition: PsiElement?, positionInFakeFile: KtElement): KtScopeContext =
+        scopeProvider.getScopeContextForPosition(this, originalPosition, positionInFakeFile)
 
     fun KtDeclaration.getSymbol(): KtSymbol = symbolProvider.getSymbol(this)
 

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.components
 
+import com.intellij.psi.PsiElement
 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
@@ -20,7 +21,11 @@ abstract class KtScopeProvider : KtAnalysisSessionComponent() {
 
     abstract fun getTypeScope(type: KtType): KtScope?
 
-    abstract fun getScopeContextForPosition(originalFile: KtFile, positionInFakeFile: KtElement): KtScopeContext
+    abstract fun getScopeContextForPosition(
+        originalFile: KtFile,
+        originalPosition: PsiElement?,
+        positionInFakeFile: KtElement
+    ): KtScopeContext
 }
 
 data class KtScopeContext(val scopes: KtCompositeScope, val implicitReceiversTypes: List<KtType>)

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

@@ -11,6 +11,8 @@ 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.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunctionCopy
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
@@ -18,6 +20,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 
 object LowLevelFirApiFacade {
@@ -64,14 +67,26 @@ object LowLevelFirApiFacade {
     fun buildCompletionContextForFunction(
         firFile: FirFile,
         element: KtNamedFunction,
+        originalElement: KtNamedFunction,
         state: FirModuleResolveState,
         phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
     ): FirCompletionContext {
         val firIdeProvider = firFile.session.firIdeProvider
+        val originalFunction = state.getOrBuildFirFor(originalElement, phase) as FirSimpleFunction
         val builtFunction = firIdeProvider.buildFunctionWithBody(element)
         val towerDataContextForStatement = mutableMapOf<FirStatement, FirTowerDataContext>()
 
-        val function = builtFunction.apply {
+        // 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
+        val frankensteinFunction = buildSimpleFunctionCopy(originalFunction) {
+            body = builtFunction.body
+            symbol = builtFunction.symbol as FirNamedFunctionSymbol
+            resolvePhase = minOf(originalFunction.resolvePhase, FirResolvePhase.DECLARATIONS)
+            this.session = state.firIdeSourcesSession
+        }
+
+        val function = frankensteinFunction.apply {
             state.lazyResolveFunctionForCompletion(this, firFile, firIdeProvider, phase, towerDataContextForStatement)
             state.recordPsiToFirMappingsForCompletionFrom(this, firFile, element.containingKtFile)
         }

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -105,14 +106,21 @@ internal class KtFirScopeProvider(
         return convertToKtScope(firTypeScope)
     }
 
-    override fun getScopeContextForPosition(originalFile: KtFile, positionInFakeFile: KtElement): KtScopeContext = withValidityAssertion {
+    override fun getScopeContextForPosition(
+        originalFile: KtFile,
+        originalPosition: PsiElement?,
+        positionInFakeFile: KtElement
+    ): KtScopeContext = withValidityAssertion {
         val originalFirFile = originalFile.getOrBuildFirOfType<FirFile>(firResolveState)
         val fakeEnclosingFunction = positionInFakeFile.getNonStrictParentOfType<KtNamedFunction>()
             ?: error("Cannot find enclosing function for ${positionInFakeFile.getElementTextInContext()}")
+        val originalEnclosingFunction = originalPosition?.getNonStrictParentOfType<KtNamedFunction>()
+            ?: error("Cannot find original enclosing function for $originalPosition")
 
         val completionContext = LowLevelFirApiFacade.buildCompletionContextForFunction(
             originalFirFile,
             fakeEnclosingFunction,
+            originalEnclosingFunction,
             state = firResolveState
         )