From b8248f01976233ea7f2f922a056ecd5ec07bdd5d Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Tue, 8 Jun 2021 13:16:38 +0200
Subject: [PATCH] FIR IDE: do not ignore parameters for extension applicability
 for completion

---
 .../api/resolver/candidateInfoProviders.kt      |  8 +++++++-
 .../KtFirCompletionCandidateChecker.kt          | 17 ++++++-----------
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/resolver/candidateInfoProviders.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/resolver/candidateInfoProviders.kt
index 753bd15fa9d..9767fd0596d 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/resolver/candidateInfoProviders.kt
+++ b/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
     }
 }
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt
index ba84b3784a3..1ffaeab7a6e 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCompletionCandidateChecker.kt
+++ b/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 ->
-            checkExtension(firFunction, originalFile, nameExpression, possibleExplicitReceiver)
+        require(firSymbolForCandidate is KtFirSymbol<*>)
+        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,