FIR IDE: do not ignore parameters for extension applicability for completion

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

@@ -7,7 +7,11 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.resolver
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.resolve.calls.*
+import org.jetbrains.kotlin.fir.resolve.inference.receiverType
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
+import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 
 /**
@@ -84,7 +88,9 @@ class CheckExtensionForCompletionCandidateInfoProvider(
     override fun shouldFailBeforeResolve(): Boolean = with(resolutionParameters) {
         val callHasExtensionReceiver = explicitReceiverKind() == ExplicitReceiverKind.EXTENSION_RECEIVER
                 || implicitExtensionReceiverValue() != null
-        val candidateHasExtensionReceiver = callableSymbol.fir.receiverTypeRef != null
+        val fir = callableSymbol.fir
+        val candidateHasExtensionReceiver = fir.receiverTypeRef != null
+                || fir is FirVariable && fir.returnTypeRef.coneType.receiverType(firSession) != null
         callHasExtensionReceiver != candidateHasExtensionReceiver
     }
 }

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

@@ -21,6 +21,8 @@ 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.KtFirKotlinPropertySymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirValueParameterSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 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
@@ -41,20 +43,13 @@ internal class KtFirCompletionCandidateChecker(
         nameExpression: KtSimpleNameExpression,
         possibleExplicitReceiver: KtExpression?,
     ): Boolean = withValidityAssertion {
-        val functionFits = firSymbolForCandidate.withResolvedFirOfType<KtFirFunctionSymbol, FirSimpleFunction, Boolean> { firFunction ->
+        require(firSymbolForCandidate is KtFirSymbol<*>)
-            checkExtension(firFunction, originalFile, nameExpression, possibleExplicitReceiver)
+        return firSymbolForCandidate.firRef.withFir(phase = FirResolvePhase.BODY_RESOLVE) { declaration ->
+            check(declaration is FirCallableDeclaration<*>)
+            checkExtension(declaration, originalFile, nameExpression, possibleExplicitReceiver)
         }
-        val propertyFits = firSymbolForCandidate.withResolvedFirOfType<KtFirKotlinPropertySymbol, FirProperty, Boolean> { firProperty ->
-            checkExtension(firProperty, originalFile, nameExpression, possibleExplicitReceiver)
-        }
-
-        functionFits ?: propertyFits ?: false
     }
 
-    private inline fun <reified T : KtFirSymbol<F>, F : FirDeclaration, R> KtCallableSymbol.withResolvedFirOfType(
-        noinline action: (F) -> R,
-    ): R? = this.safeAs<T>()?.firRef?.withFir(phase = FirResolvePhase.BODY_RESOLVE, action)
-
     private fun checkExtension(
         candidateSymbol: FirCallableDeclaration<*>,
         originalFile: KtFile,