K2: Propagate annotation of function type kind in compose plugin

We want to test whether the propagation of function type kind annotation
correctly works or not. This commit updates
ComposableFunctionsTransformer that mimics the real compose compiler
plugin to propagate the composable annotation i.e., @MyComposable from
MyComposableFunction kind. Finally, the generated IR function for lambda
expression must have the composable annotation, when the lambda
expression does not have the explicit composable annotation, but we
infer the composable annotation based on the annotation of function type
for the value parameter.

^KT-64994

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/ComposableFunctionsTransformer.kt

@@ -6,22 +6,28 @@
 package org.jetbrains.kotlin.ir.plugin
 
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
-import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.builtins.functions.FunctionTypeKind
 import org.jetbrains.kotlin.fir.plugin.types.ComposableNames.FULL_COMPOSABLE_NAME_PREFIX
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrValueParameter
+import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
+import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.impl.IrTypeBase
 import org.jetbrains.kotlin.ir.types.impl.IrTypeProjectionImpl
-import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
-import org.jetbrains.kotlin.ir.util.functions
+import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
 class ComposableFunctionsTransformer(val pluginContext: IrPluginContext) : IrElementVisitorVoid {
@@ -66,6 +72,45 @@ class ComposableFunctionsTransformer(val pluginContext: IrPluginContext) : IrEle
         visitElement(expression)
     }
 
+    /**
+     * This function propagates the special function type kind for composable to function expressions like lambda expression.
+     */
+    override fun visitFunctionExpression(expression: IrFunctionExpression) {
+        if (expression.type.isSyntheticComposableFunction()) {
+            expression.function.mark()
+        }
+        super.visitFunctionExpression(expression)
+    }
+
+    /**
+     * This function propagates the special function type kind for composable to function references.
+     */
+    override fun visitFunctionReference(expression: IrFunctionReference) {
+        if (expression.type.isSyntheticComposableFunction()) {
+            expression.symbol.owner.mark()
+        }
+        super.visitFunctionReference(expression)
+    }
+
+    private fun IrType.isSyntheticComposableFunction() =
+        classOrNull?.owner?.let {
+            it.name.asString().startsWith("MyComposableFunction") &&
+                    it.packageFqName?.asString() == "some"
+        } ?: false
+
+    private val composableClassId = ClassId(FqName("org.jetbrains.kotlin.fir.plugin"), FqName("MyComposable"), false)
+
+    private val composableSymbol = pluginContext.referenceClass(composableClassId)!!
+
+    private fun IrFunction.mark() {
+        if (!hasAnnotation(composableClassId)) {
+            annotations = annotations + IrConstructorCallImpl.fromSymbolOwner(
+                composableSymbol.owner.defaultType,
+                composableSymbol.constructors.single(),
+            )
+        }
+    }
+
     private fun updateReferenceInCallIfNeeded(call: IrCall) {
         val function = call.symbol.owner
         if (function.name != INVOKE) return

plugins/fir-plugin-prototype/testData/box/composableFunction.fir.ir.txt

@@ -64,6 +64,8 @@ FILE fqName:<root> fileName:/composableFunction.kt
       VAR name:l1 type:kotlin.Function0<kotlin.Unit> [val]
         FUN_EXPR type=kotlin.Function0<kotlin.Unit> origin=LAMBDA
           FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.Unit
+            annotations:
+              MyComposable
             BLOCK_BODY
               TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
                 BLOCK type=kotlin.Int origin=POSTFIX_INCR
@@ -77,6 +79,8 @@ FILE fqName:<root> fileName:/composableFunction.kt
       VAR name:l2 type:@[MyComposable] kotlin.Function0<kotlin.Unit> [val]
         FUN_EXPR type=kotlin.Function0<kotlin.Unit> origin=LAMBDA
           FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.Unit
+            annotations:
+              MyComposable
             BLOCK_BODY
               TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
                 BLOCK type=kotlin.Int origin=POSTFIX_INCR
@@ -126,6 +130,8 @@ FILE fqName:<root> fileName:/composableFunction.kt
       CALL 'public final fun runComposable (block: @[MyComposable] kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
         block: FUN_EXPR type=kotlin.Function0<kotlin.Unit> origin=LAMBDA
           FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.Unit
+            annotations:
+              MyComposable
             BLOCK_BODY
               TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
                 BLOCK type=kotlin.Int origin=POSTFIX_INCR
@@ -157,6 +163,8 @@ FILE fqName:<root> fileName:/composableFunction.kt
         CALL 'public final fun runComposable2 (block: kotlin.Function1<kotlin.String, kotlin.String>): kotlin.String declared in <root>' type=kotlin.String origin=null
           block: FUN_EXPR type=kotlin.Function1<kotlin.String, kotlin.String> origin=LAMBDA
             FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.String) returnType:kotlin.String
+              annotations:
+                MyComposable
               VALUE_PARAMETER name:it index:0 type:kotlin.String
               BLOCK_BODY
                 RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.String): kotlin.String declared in <root>.test_2'