FIR Completion: Drop originalPosition parameter from getScopeContextForPosition

- It is easier to retrieve the original enclosing function by the fake
enclosing function than trying to correct the `originalPosition` in
different situations

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

@@ -67,35 +67,6 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
 private object KotlinAvailableScopesCompletionProvider {
     private val lookupElementFactory = KotlinFirLookupElementFactory()
 
-    fun getOriginalPosition(parameters: CompletionParameters, originalFile: KtFile): PsiElement {
-        fun PsiElement.getPositionForExpressionFunctionBody(): PsiElement? {
-            return when {
-                this is KtNamedFunction && bodyExpression == null && equalsToken != null -> this
-                elementType == KtTokens.EQ && parent is KtNamedFunction -> this
-                else -> null
-            }
-        }
-
-        val originalPosition = parameters.originalPosition
-        if (originalPosition is PsiWhiteSpace) {
-            originalPosition.prevSibling?.getPositionForExpressionFunctionBody()?.let {
-                /* We are in expression function body
-                 * fun x() = <caret>
-                 */
-                return it
-            }
-        }
-        if (originalPosition != null) return originalPosition
-        if (parameters.offset == originalFile.textLength) {
-            /*
-             * We are in the end of the file possibly in function with expression body
-             * fun x() = <caret><EOF>
-             */
-            originalFile.findElementAt(parameters.offset - 1)?.getPositionForExpressionFunctionBody()?.let { return it }
-        }
-        error("Can not find original position")
-    }
-
     @OptIn(InvalidWayOfUsingAnalysisSession::class)
     fun addCompletions(parameters: CompletionParameters, result: CompletionResultSet) {
         val originalFile = parameters.originalFile as? KtFile ?: return
@@ -105,10 +76,8 @@ private object KotlinAvailableScopesCompletionProvider {
 
         val possibleReceiver = nameExpression.getQualifiedExpressionForSelector()?.receiverExpression
 
-        val originalPosition = getOriginalPosition(parameters, originalFile)
-
         with(getAnalysisSessionFor(originalFile).createContextDependentCopy()) {
-            val (implicitScopes, implicitReceivers) = originalFile.getScopeContextForPosition(originalPosition, nameExpression)
+            val (implicitScopes, implicitReceivers) = originalFile.getScopeContextForPosition(nameExpression)
 
             val typeOfPossibleReceiver = possibleReceiver?.getKtType()
             val possibleReceiverScope = typeOfPossibleReceiver?.getTypeScope()

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

@@ -69,8 +69,8 @@ abstract class KtAnalysisSession(override val token: ValidityToken) : ValidityTo
 
     fun KtType.getTypeScope(): KtScope? = scopeProvider.getTypeScope(this)
 
-    fun KtFile.getScopeContextForPosition(originalPosition: PsiElement, positionInFakeFile: KtElement): KtScopeContext =
+    fun KtFile.getScopeContextForPosition(positionInFakeFile: KtElement): KtScopeContext =
-        scopeProvider.getScopeContextForPosition(this, originalPosition, positionInFakeFile)
+        scopeProvider.getScopeContextForPosition(this, positionInFakeFile)
 
     fun KtDeclaration.getSymbol(): KtSymbol = symbolProvider.getSymbol(this)
 

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

@@ -23,7 +23,6 @@ abstract class KtScopeProvider : KtAnalysisSessionComponent() {
 
     abstract fun getScopeContextForPosition(
         originalFile: KtFile,
-        originalPosition: PsiElement,
         positionInFakeFile: KtElement
     ): KtScopeContext
 }

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

@@ -39,8 +39,6 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymb
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
 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.KtCallableSymbol
-import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassLikeSymbol
 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
@@ -111,14 +109,13 @@ internal class KtFirScopeProvider(
 
     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>()
+        val originalEnclosingFunction = originalFile.findFunctionDeclarationAt(fakeEnclosingFunction.textOffset)
-            ?: error("Cannot find original enclosing function for $originalPosition")
+            ?: error("Cannot find original function matching to ${fakeEnclosingFunction.getElementTextInContext()} in $originalFile")
 
         val completionContext = LowLevelFirApiFacade.buildCompletionContextForFunction(
             originalFirFile,
@@ -146,6 +143,11 @@ internal class KtFirScopeProvider(
         KtScopeContext(getCompositeScope(allKtScopes), implicitReceiversTypes)
     }
 
+    private fun KtFile.findFunctionDeclarationAt(offset: Int): KtNamedFunction? =
+        findElementAt(offset)
+            ?.getNonStrictParentOfType<KtNamedFunction>()
+            ?.takeIf { it.textOffset == offset }
+
     private fun convertToKtScope(firScope: FirScope): KtScope {
         firScopeStorage.register(firScope)
         return when (firScope) {