FIR IDE: output substitutor from checkExtensionIsSuitable

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirAnalysisSessionComponent.kt

@@ -5,17 +5,6 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.components
 
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.FirSourceElement
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
-import org.jetbrains.kotlin.fir.analysis.diagnostics.toFirDiagnostics
-import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
-import org.jetbrains.kotlin.fir.typeContext
-import org.jetbrains.kotlin.fir.types.ConeInferenceContext
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.ConeStarProjection
-import org.jetbrains.kotlin.fir.types.ConeTypeProjection
 import org.jetbrains.kotlin.analysis.api.KtStarProjectionTypeArgument
 import org.jetbrains.kotlin.analysis.api.KtTypeArgument
 import org.jetbrains.kotlin.analysis.api.KtTypeArgumentWithVariance
@@ -23,9 +12,25 @@ import org.jetbrains.kotlin.analysis.api.diagnostics.KtDiagnosticWithPsi
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.diagnostics.KT_DIAGNOSTIC_CONVERTER
 import org.jetbrains.kotlin.analysis.api.fir.types.KtFirType
+import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
+import org.jetbrains.kotlin.fir.analysis.diagnostics.toFirDiagnostics
+import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
+import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
+import org.jetbrains.kotlin.fir.typeContext
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.types.TypeCheckerState
 import org.jetbrains.kotlin.types.model.convertVariance
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal interface KtFirAnalysisSessionComponent {
     val analysisSession: KtFirAnalysisSession
@@ -67,6 +72,30 @@ internal interface KtFirAnalysisSessionComponent {
 
     fun createTypeCheckerContext(): TypeCheckerState {
         // TODO use correct session here,
-        return analysisSession.firResolveState.rootModuleSession.typeContext.newTypeCheckerState(errorTypesEqualToAnything = true, stubTypesEqualToAnything = true)
+        return analysisSession.firResolveState.rootModuleSession.typeContext.newTypeCheckerState(
+            errorTypesEqualToAnything = true,
+            stubTypesEqualToAnything = true
+        )
+    }
+
+    fun FirQualifiedAccessExpression.createSubstitutorFromTypeArguments(): KtSubstitutor? {
+        val symbol = when (val calleeReference = calleeReference) {
+            is FirResolvedNamedReference -> calleeReference.resolvedSymbol as? FirCallableSymbol<*>
+            is FirErrorNamedReference -> calleeReference.candidateSymbol as? FirCallableSymbol<*>
+            else -> null
+        } ?: return null
+        return createSubstitutorFromTypeArguments(symbol)
+    }
+
+    fun FirQualifiedAccessExpression.createSubstitutorFromTypeArguments(functionSymbol: FirCallableSymbol<*>): KtSubstitutor {
+        val typeArgumentMap = mutableMapOf<FirTypeParameterSymbol, ConeKotlinType>()
+        for (i in typeArguments.indices) {
+            val type = typeArguments[i].safeAs<FirTypeProjectionWithVariance>()?.typeRef?.coneType
+            if (type != null) {
+                typeArgumentMap[functionSymbol.typeParameterSymbols[i]] = type
+            }
+        }
+        val coneSubstitutor = substitutorByMap(typeArgumentMap, rootModuleSession)
+        return firSymbolBuilder.typeBuilder.buildSubstitutor(coneSubstitutor)
     }
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt

@@ -36,10 +36,8 @@ import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.impl.FirSimpleNamedReference
 import org.jetbrains.kotlin.fir.resolve.calls.FirErrorReferenceWithCandidate
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper
 import org.jetbrains.kotlin.idea.references.readWriteAccess
@@ -226,26 +224,8 @@ internal class KtFirCallResolver(
     }
 
     private fun FirFunctionCall.asSimpleFunctionCall(): KtFunctionCall? {
-        val calleeReference = this.calleeReference
+        val target = this.calleeReference.createCallTarget() ?: return null
-        val target = calleeReference.createCallTarget() ?: return null
+        return KtFunctionCall(createArgumentMapping(), target, createSubstitutorFromTypeArguments() ?: return null, token)
-        val symbol = when (calleeReference) {
-            is FirResolvedNamedReference -> calleeReference.resolvedSymbol as? FirCallableSymbol<*>
-            is FirErrorNamedReference -> calleeReference.candidateSymbol as? FirCallableSymbol<*>
-            else -> null
-        } ?: return null
-        return KtFunctionCall(createArgumentMapping(), target, createSubstitutorFromTypeArguments(symbol), token)
-    }
-
-    private fun FirFunctionCall.createSubstitutorFromTypeArguments(functionSymbol: FirCallableSymbol<*>): KtSubstitutor {
-        val typeArgumentMap = mutableMapOf<FirTypeParameterSymbol, ConeKotlinType>()
-        for (i in typeArguments.indices) {
-            val type = typeArguments[i].safeAs<FirTypeProjectionWithVariance>()?.typeRef?.coneType
-            if (type != null) {
-                typeArgumentMap[functionSymbol.typeParameterSymbols[i]] = type
-            }
-        }
-        val coneSubstitutor = substitutorByMap(typeArgumentMap, rootModuleSession)
-        return firSymbolBuilder.typeBuilder.buildSubstitutor(coneSubstitutor)
     }
 
     private fun FirAnnotationCall.asAnnotationCall(): KtAnnotationCall? {

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

@@ -5,20 +5,6 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.components
 
-import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.declarations.FirVariable
-import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
-import org.jetbrains.kotlin.fir.resolve.inference.receiverType
-import org.jetbrains.kotlin.fir.types.coneType
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir.getTowerContextProvider
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
-import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.ResolutionParameters
-import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.SingleCandidateResolutionMode
-import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.SingleCandidateResolver
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getElementTextInContext
 import org.jetbrains.kotlin.analysis.api.components.KtCompletionCandidateChecker
 import org.jetbrains.kotlin.analysis.api.components.KtExtensionApplicabilityResult
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
@@ -26,7 +12,22 @@ import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.analysis.api.fir.utils.weakRef
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
+import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
 import org.jetbrains.kotlin.analysis.api.withValidityAssertion
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir.getTowerContextProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.ResolutionParameters
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.SingleCandidateResolutionMode
+import org.jetbrains.kotlin.analysis.low.level.api.fir.resolver.SingleCandidateResolver
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getElementTextInContext
+import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.FirVariable
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.fir.resolve.inference.receiverType
+import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
@@ -70,17 +71,18 @@ internal class KtFirCompletionCandidateChecker(
                 explicitReceiver = explicitReceiverExpression
             )
             resolver.resolveSingleCandidate(resolutionParameters)?.let {
+                val substitutor = it.createSubstitutorFromTypeArguments() ?: return@let null
                 return when {
                     candidateSymbol is FirVariable && candidateSymbol.returnTypeRef.coneType.receiverType(rootModuleSession) != null -> {
-                        KtExtensionApplicabilityResult.ApplicableAsFunctionalVariableCall
+                        KtExtensionApplicabilityResult.ApplicableAsFunctionalVariableCall(substitutor)
                     }
                     else -> {
-                        KtExtensionApplicabilityResult.ApplicableAsExtensionCallable
+                        KtExtensionApplicabilityResult.ApplicableAsExtensionCallable(substitutor)
                     }
                 }
             }
         }
-        return KtExtensionApplicabilityResult.NonApplicable
+        return KtExtensionApplicabilityResult.NonApplicable(KtSubstitutor.Empty(token))
     }
 
     private fun getImplicitReceivers(fakeNameExpression: KtSimpleNameExpression): Sequence<ImplicitReceiverValue<*>?> {

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtCompletionCandidateChecker.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.api.components
 
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
@@ -19,10 +20,21 @@ public abstract class KtCompletionCandidateChecker : KtAnalysisSessionComponent(
     ): KtExtensionApplicabilityResult
 }
 
-public enum class KtExtensionApplicabilityResult(public val isApplicable: Boolean) {
+public sealed class KtExtensionApplicabilityResult {
-    ApplicableAsExtensionCallable(isApplicable = true),
+    public abstract val isApplicable: Boolean
-    ApplicableAsFunctionalVariableCall(isApplicable = true),
+    public abstract val substitutor: KtSubstitutor
-    NonApplicable(isApplicable = false),
+
+    public class ApplicableAsExtensionCallable(override val substitutor: KtSubstitutor) : KtExtensionApplicabilityResult() {
+        override val isApplicable: Boolean get() = true
+    }
+
+    public class ApplicableAsFunctionalVariableCall(override val substitutor: KtSubstitutor) : KtExtensionApplicabilityResult() {
+        override val isApplicable: Boolean get() = true
+    }
+
+    public class NonApplicable(override val substitutor: KtSubstitutor) : KtExtensionApplicabilityResult() {
+        override val isApplicable: Boolean = false
+    }
 }
 
 public interface KtCompletionCandidateCheckerMixIn : KtAnalysisSessionMixIn {