Repository: JetBrains/kotlin
File: plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunInterfaceLowering.kt (Kotlin)

Code before:
import org.jetbrains.kotlin.descriptors.Modality
import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
import org.jetbrains.kotlin.ir.expressions.IrTypeOperator.IMPLICIT_CAST
import org.jetbrains.kotlin.ir.expressions.IrTypeOperator.SAM_CONVERSION
import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
import org.jetbrains.kotlin.ir.expressions.impl.IrRichFunctionReferenceImpl
import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
import org.jetbrains.kotlin.ir.types.IrType
import org.jetbrains.kotlin.ir.types.IrTypeSystemContextImpl
import org.jetbrains.kotlin.ir.types.classOrNull
import org.jetbrains.kotlin.ir.util.functions
import org.jetbrains.kotlin.ir.util.isLambda
import org.jetbrains.kotlin.ir.util.selectSAMOverriddenFunction
import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
import org.jetbrains.kotlin.platform.jvm.isJvm
import org.jetbrains.kotlin.platform.konan.isNative

@Suppress("PRE_RELEASE_CLASS")
class ComposableFunInterfaceLowering(private val context: IrPluginContext) :
    IrElementTransformerVoidWithContext(),
    ModuleLoweringPass {

    override fun lower(irModule: IrModuleFragment) {
        if (context.platform.isJvm() || context.platform.isNative()) {
            irModule.transformChildrenVoid(this)
        }
    }

    override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
        val functionExpr = expression.findSamFunctionExpr()
        if (functionExpr != null && expression.typeOperand.isComposableFunInterface()) {
            val argument = functionExpr.transform(this, null) as IrFunctionExpression
            val superType = expression.typeOperand
            val superClass = superType.classOrNull ?: error("Expected non-null class")
            return if (!context.platform.isJvm()) {
                IrRichFunctionReferenceImpl(
                    startOffset = expression.startOffset,
                    endOffset = expression.endOffset,
                    type = expression.typeOperand,
                    reflectionTargetSymbol = null,
                    overriddenFunctionSymbol = superClass.owner.selectSAMOverriddenFunction().symbol,
                    invokeFunction = functionExpr.function,
                    origin = functionExpr.origin
                )
            } else {
                FunctionReferenceBuilder(
                    argument,
                    superClass,
                    superType,

Review comment (suggestion):
Can you please also check 

`pluginContext.languageVersionSettings.supportsFeature(LanguageFeature.IrRichCallableReferencesInKlibs)` here?

Otherwise, it would break with export to 2.2 after being backported to 2.3. In 2.4 branch, this check is not needed anymore (as you can't have 2.4 with LV 2.2 outside of jvm), so if you prefer to add it only to 2.3.20 branch, that's also an option. 

Code after:
import org.jetbrains.kotlin.config.LanguageFeature
import org.jetbrains.kotlin.descriptors.Modality
import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
import org.jetbrains.kotlin.ir.expressions.IrTypeOperator.IMPLICIT_CAST
import org.jetbrains.kotlin.ir.expressions.IrTypeOperator.SAM_CONVERSION
import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
import org.jetbrains.kotlin.ir.expressions.impl.IrRichFunctionReferenceImpl
import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
import org.jetbrains.kotlin.ir.types.IrType
import org.jetbrains.kotlin.ir.types.IrTypeSystemContextImpl
import org.jetbrains.kotlin.ir.types.classOrNull
import org.jetbrains.kotlin.ir.util.functions
import org.jetbrains.kotlin.ir.util.isLambda
import org.jetbrains.kotlin.ir.util.selectSAMOverriddenFunction
import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
import org.jetbrains.kotlin.platform.jvm.isJvm
import org.jetbrains.kotlin.platform.konan.isNative

@Suppress("PRE_RELEASE_CLASS")
class ComposableFunInterfaceLowering(private val context: IrPluginContext) :
    IrElementTransformerVoidWithContext(),
    ModuleLoweringPass {

    override fun lower(irModule: IrModuleFragment) {
        if (context.platform.isJvm() || (context.platform.isNative() && context.languageVersionSettings.supportsFeature(LanguageFeature.IrRichCallableReferencesInKlibs))) {
            irModule.transformChildrenVoid(this)
        }
    }

    override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
        val functionExpr = expression.findSamFunctionExpr()
        if (functionExpr != null && expression.typeOperand.isComposableFunInterface()) {
            val argument = functionExpr.transform(this, null) as IrFunctionExpression
            val superType = expression.typeOperand
            val superClass = superType.classOrNull ?: error("Expected non-null class")
            return if (!context.platform.isJvm()) {
                IrRichFunctionReferenceImpl(
                    startOffset = expression.startOffset,
                    endOffset = expression.endOffset,
                    type = expression.typeOperand,
                    reflectionTargetSymbol = null,
                    overriddenFunctionSymbol = superClass.owner.selectSAMOverriddenFunction().symbol,
                    invokeFunction = functionExpr.function,
                    origin = functionExpr.origin
                )
            } else {
                FunctionReferenceBuilder(
                    argument,
                    superClass,