[FIR] Add extension for injecting new implicit receivers after resolution of function call

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt

@@ -28,6 +28,7 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirAdditionalCheckersExtension::class,
             FirSupertypeGenerationExtension::class,
             FirTypeAttributeExtension::class,
+            FirExpressionResolutionExtension::class,
         )
     }
 
@@ -61,6 +62,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
         operator fun ((FirSession) -> FirTypeAttributeExtension).unaryPlus() {
             registerExtension(FirTypeAttributeExtension::class, FirTypeAttributeExtension.Factory { this.invoke(it) })
         }
+
+        @JvmName("plusExpressionResolutionExtension")
+        operator fun ((FirSession) -> FirExpressionResolutionExtension).unaryPlus() {
+            registerExtension(FirExpressionResolutionExtension::class, FirExpressionResolutionExtension.Factory { this.invoke(it) })
+        }
     }
 
     @OptIn(PluginServicesInitialization::class)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExpressionResolutionExtension.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import kotlin.reflect.KClass
+
+abstract class FirExpressionResolutionExtension(session: FirSession) : FirExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("ExpressionResolutionExtension")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension> = FirExpressionResolutionExtension::class
+
+    abstract fun addNewImplicitReceivers(functionCall: FirFunctionCall): List<ConeKotlinType>
+
+    fun interface Factory : FirExtension.Factory<FirExpressionResolutionExtension>
+}
+
+val FirExtensionService.expressionResolutionExtensions: List<FirExpressionResolutionExtension> by FirExtensionService.registeredExtensions()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -16,6 +16,8 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
 import org.jetbrains.kotlin.fir.expressions.impl.toAnnotationArgumentMapping
+import org.jetbrains.kotlin.fir.extensions.expressionResolutionExtensions
+import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildExplicitSuperReference
@@ -32,6 +34,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.StoreReceiver
 import org.jetbrains.kotlin.fir.scopes.impl.isWrappedIntegerOperator
 import org.jetbrains.kotlin.fir.scopes.impl.isWrappedIntegerOperatorForUnsignedType
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
@@ -56,6 +59,8 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
     var enableArrayOfCallTransformation = false
     var containingSafeCallExpression: FirSafeCallExpression? = null
 
+    private val expressionResolutionExtensions = session.extensionService.expressionResolutionExtensions.takeIf { it.isNotEmpty() }
+
     init {
         @Suppress("LeakingThis")
         components.callResolver.initTransformer(this)
@@ -378,6 +383,9 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             }
         val result = completeInference.transformToIntegerOperatorCallOrApproximateItIfNeeded(data)
         dataFlowAnalyzer.exitFunctionCall(result, callCompleted)
+
+        addReceiversFromExtensions(result)
+
         if (callCompleted) {
             if (enableArrayOfCallTransformation) {
                 return arrayOfCallTransformer.transformFunctionCall(result, null)
@@ -386,6 +394,23 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         return result
     }
 
+    @OptIn(PrivateForInline::class)
+    private fun addReceiversFromExtensions(functionCall: FirFunctionCall) {
+        val extensions = expressionResolutionExtensions ?: return
+        val boundSymbol = context.containerIfAny?.symbol as? FirCallableSymbol<*> ?: return
+        for (extension in extensions) {
+            for (receiverType in extension.addNewImplicitReceivers(functionCall)) {
+                val receiverValue = ImplicitExtensionReceiverValue(
+                    boundSymbol,
+                    receiverType,
+                    session,
+                    scopeSession
+                )
+                context.addReceiver(name = null, receiverValue)
+            }
+        }
+    }
+
     private fun FirFunctionCall.transformToIntegerOperatorCallOrApproximateItIfNeeded(resolutionMode: ResolutionMode): FirFunctionCall {
         if (!explicitReceiver.isIntegerLiteralOrOperatorCall()) return this
         val resolvedSymbol = when (val reference = calleeReference) {