[FIR] Introduce extension for custom SAM conversions

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

@@ -9,6 +9,7 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
+import org.jetbrains.kotlin.fir.resolve.FirSamConversionTransformerExtension
 import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.reflect.KClass
 
@@ -31,6 +32,7 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirTypeAttributeExtension::class,
             FirExpressionResolutionExtension::class,
             FirExtensionSessionComponent::class,
+            FirSamConversionTransformerExtension::class,
         )
     }
 
@@ -74,6 +76,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             registerExtension(FirExtensionSessionComponent::class, this)
         }
 
+        @JvmName("plusSamConversionTransformerExtension")
+        operator fun (FirSamConversionTransformerExtension.Factory).unaryPlus() {
+            registerExtension(FirSamConversionTransformerExtension::class, this)
+        }
+
         // ------------------ reference methods ------------------
 
         @JvmName("plusStatusTransformerExtension")
@@ -110,6 +117,23 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
         operator fun ((FirSession) -> FirExtensionSessionComponent).unaryPlus() {
             FirExtensionSessionComponent.Factory { this.invoke(it) }.unaryPlus()
         }
+
+        @JvmName("plusSamConversionTransformerExtension")
+        operator fun ((FirSession) -> FirSamConversionTransformerExtension).unaryPlus() {
+            FirSamConversionTransformerExtension.Factory { this.invoke(it) }.unaryPlus()
+        }
+
+        // ------------------ utilities ------------------
+
+        @JvmName("bindLeft")
+        fun <T, R> ((T, FirSession) -> R).bind(value: T): (FirSession) -> R {
+            return { this.invoke(value, it) }
+        }
+
+        @JvmName("bindRight")
+        fun <T, R> ((FirSession, T) -> R).bind(value: T): (FirSession) -> R {
+            return { this.invoke(it, value) }
+        }
     }
 
     @OptIn(PluginServicesInitialization::class)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirSamConversionTransformerExtension.kt

@@ -0,0 +1,31 @@
+/*
+ * 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.resolve
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.extensions.FirExtension
+import org.jetbrains.kotlin.fir.extensions.FirExtensionPointName
+import org.jetbrains.kotlin.fir.extensions.FirExtensionService
+import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
+import kotlin.reflect.KClass
+
+abstract class FirSamConversionTransformerExtension(session: FirSession) : FirExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("SamConversionTransformerExtension")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension> = FirSamConversionTransformerExtension::class
+
+    abstract fun getCustomFunctionalTypeForSamConversion(function: FirSimpleFunction): ConeLookupTagBasedType?
+
+    fun interface Factory : FirExtension.Factory<FirSamConversionTransformerExtension>
+}
+
+val FirExtensionService.samConversionTransformers: List<FirSamConversionTransformerExtension> by FirExtensionService.registeredExtensions()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SamResolution.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.declarations.utils.*
 import org.jetbrains.kotlin.fir.diagnostics.ConeIntermediateDiagnostic
+import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticFunctionSymbol
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
@@ -58,6 +59,7 @@ class FirSamResolverImpl(
 ) : FirSamResolver() {
     private val resolvedFunctionType: NullableMap<FirRegularClass, SAMInfo<ConeLookupTagBasedType>?> = NullableMap()
     private val samConstructorsCache = session.samConstructorStorage.samConstructors
+    private val samConversionTransformers = session.extensionService.samConversionTransformers
 
     override fun getSamInfoForPossibleSamType(type: ConeKotlinType): SAMInfo<ConeKotlinType>? {
         return when (type) {
@@ -243,7 +245,11 @@ class FirSamResolverImpl(
             val abstractMethod = firRegularClass.getSingleAbstractMethodOrNull(session, scopeSession) ?: return@getOrPut null
             // TODO: val shouldConvertFirstParameterToDescriptor = samWithReceiverResolvers.any { it.shouldConvertFirstSamParameterToReceiver(abstractMethod) }
 
-            SAMInfo(abstractMethod.symbol, abstractMethod.getFunctionTypeForAbstractMethod())
+            val typeFromExtension = samConversionTransformers.firstNotNullOfOrNull {
+                it.getCustomFunctionalTypeForSamConversion(abstractMethod)
+            }
+
+            SAMInfo(abstractMethod.symbol, typeFromExtension ?: abstractMethod.getFunctionTypeForAbstractMethod())
         }
     }