[FIR plugin] Replace composable functions with regular kotlin functions

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/types/ComposableLikeFunctionTypeKindExtension.kt

@@ -20,13 +20,19 @@ class ComposableLikeFunctionTypeKindExtension(session: FirSession) : FirFunction
     }
 }
 
-private val COMPOSABLE_PACKAGE_FQN = FqName.topLevel(Name.identifier("some"))
+object ComposableNames {
-private val MY_COMPOSABLE_ANNOTATION_CLASS_ID = ClassId.topLevel("MyComposable".fqn())
+    val COMPOSABLE_PACKAGE_FQN = FqName.topLevel(Name.identifier("some"))
+    val MY_COMPOSABLE_ANNOTATION_CLASS_ID = ClassId.topLevel("MyComposable".fqn())
+    val COMPOSABLE_NAME_PREFIX = "MyComposableFunction"
+    val KCOMPOSABLE_NAME_PREFIX = "KMyComposableFunction"
+
+    val FULL_COMPOSABLE_NAME_PREFIX = COMPOSABLE_PACKAGE_FQN.child(Name.identifier(COMPOSABLE_NAME_PREFIX)).asString()
+}
 
 object ComposableFunction : FunctionTypeKind(
-    COMPOSABLE_PACKAGE_FQN,
+    ComposableNames.COMPOSABLE_PACKAGE_FQN,
-    "MyComposableFunction",
+    ComposableNames.COMPOSABLE_NAME_PREFIX,
-    MY_COMPOSABLE_ANNOTATION_CLASS_ID,
+    ComposableNames.MY_COMPOSABLE_ANNOTATION_CLASS_ID,
     isReflectType = false
 ) {
     override val prefixForTypeRender: String
@@ -39,9 +45,9 @@ object ComposableFunction : FunctionTypeKind(
 }
 
 object KComposableFunction : FunctionTypeKind(
-    COMPOSABLE_PACKAGE_FQN,
+    ComposableNames.COMPOSABLE_PACKAGE_FQN,
-    "KMyComposableFunction",
+    ComposableNames.KCOMPOSABLE_NAME_PREFIX,
-    MY_COMPOSABLE_ANNOTATION_CLASS_ID,
+    ComposableNames.MY_COMPOSABLE_ANNOTATION_CLASS_ID,
     isReflectType = true
 ) {
     override val serializeAsFunctionWithAnnotationUntil: String

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

@@ -0,0 +1,123 @@
+/*
+ * Copyright 2010-2023 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.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.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.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+
+class ComposableFunctionsTransformer(val pluginContext: IrPluginContext) : IrElementVisitorVoid {
+    companion object {
+        private val INVOKE = Name.identifier("invoke")
+    }
+
+    override fun visitElement(element: IrElement) {
+        element.acceptChildrenVoid(this)
+    }
+
+    override fun visitValueParameter(declaration: IrValueParameter) {
+        declaration.apply {
+            type = type.update()
+        }
+        visitElement(declaration)
+    }
+
+    override fun visitFunction(declaration: IrFunction) {
+        declaration.apply {
+            returnType = returnType.update()
+        }
+        visitElement(declaration)
+    }
+
+    override fun visitCall(expression: IrCall) {
+        updateReferenceInCallIfNeeded(expression)
+        visitElement(expression)
+    }
+
+    override fun visitVariable(declaration: IrVariable) {
+        declaration.apply {
+            type = type.update()
+        }
+        visitElement(declaration)
+    }
+
+    override fun visitExpression(expression: IrExpression) {
+        expression.apply {
+            type = type.update()
+        }
+        visitElement(expression)
+    }
+
+    private fun updateReferenceInCallIfNeeded(call: IrCall) {
+        val function = call.symbol.owner
+        if (function.name != INVOKE) return
+        val functionClass = function.parent as? IrClass ?: return
+        val updatedClass = calculateUpdatedClass(functionClass) ?: return
+        val newFunction = updatedClass.functions.firstOrNull { it.name == INVOKE } ?: return
+        call.symbol = newFunction.symbol
+    }
+
+    private fun IrType.update(): IrType {
+        return substitutor.substitute(this)
+    }
+
+    private val substitutor = TypeSubstitutor()
+
+    private inner class TypeSubstitutor : AbstractIrTypeSubstitutor() {
+        override fun substitute(type: IrType): IrType {
+            if (type !is IrSimpleTypeImpl) return type
+            val newArguments = type.arguments.map { substituteArgument(it) }
+            val newClassifier = calculateUpdatedClassifier(type.classifier) ?: type.classifier
+            return IrSimpleTypeImpl(
+                newClassifier,
+                type.nullability,
+                newArguments,
+                type.annotations,
+                type.abbreviation
+            )
+        }
+
+        private fun substituteArgument(typeArgument: IrTypeArgument): IrTypeArgument {
+            return when (typeArgument) {
+                is IrStarProjection -> typeArgument
+                is IrTypeProjectionImpl -> typeArgument
+                is IrTypeBase -> substitute(typeArgument) as IrTypeBase
+                else -> error("Unexpected type argument: $typeArgument")
+            }
+        }
+    }
+
+    private fun calculateUpdatedClass(irClass: IrClass): IrClass? {
+        return calculateUpdatedClassifier(irClass.symbol)?.owner as? IrClass
+    }
+
+    private fun calculateUpdatedClassifier(classifier: IrClassifierSymbol): IrClassifierSymbol? {
+        val irClass = classifier.owner as? IrClass ?: return null
+        val fqName = irClass.fqNameWhenAvailable ?: return null
+        val fqNameString = fqName.asString()
+        if (!fqNameString.startsWith(FULL_COMPOSABLE_NAME_PREFIX)) return null
+        val number = fqNameString.removePrefix(FULL_COMPOSABLE_NAME_PREFIX).toIntOrNull() ?: return null
+        val builtinClassId = FunctionTypeKind.Function.run {
+            ClassId(packageFqName, Name.identifier("$classNamePrefix$number"))
+        }
+        return pluginContext.referenceClass(builtinClassId)
+    }
+}

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

@@ -19,6 +19,7 @@ class GeneratedDeclarationsIrBodyFiller : IrGenerationExtension {
             TransformerForTopLevelDeclarationsGenerator(pluginContext),
             AllPropertiesConstructorIrGenerator(pluginContext),
             TransformerForAddingAnnotations(pluginContext),
+            ComposableFunctionsTransformer(pluginContext),
         )
 
         for (transformer in transformers) {

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

@@ -0,0 +1,165 @@
+FILE fqName:<root> fileName:/composableFunction.kt
+  FUN name:runUsual visibility:public modality:FINAL <> (block:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit
+    VALUE_PARAMETER name:block index:0 type:kotlin.Function0<kotlin.Unit>
+    BLOCK_BODY
+      CALL 'public abstract fun invoke (): R of kotlin.Function0 declared in kotlin.Function0' type=kotlin.Unit origin=INVOKE
+        $this: GET_VAR 'block: kotlin.Function0<kotlin.Unit> declared in <root>.runUsual' type=kotlin.Function0<kotlin.Unit> origin=null
+  FUN name:runComposable visibility:public modality:FINAL <> (block:@[MyComposable] kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit
+    VALUE_PARAMETER name:block index:0 type:@[MyComposable] kotlin.Function0<kotlin.Unit>
+    BLOCK_BODY
+      CALL 'public abstract fun invoke (): R of kotlin.Function0 declared in kotlin.Function0' type=kotlin.Unit origin=null
+        $this: GET_VAR 'block: @[MyComposable] kotlin.Function0<kotlin.Unit> declared in <root>.runComposable' type=@[MyComposable] kotlin.Function0<kotlin.Unit> origin=null
+  FUN name:test_1 visibility:public modality:FINAL <> () returnType:kotlin.Int
+    BLOCK_BODY
+      VAR name:x type:kotlin.Int [var]
+        CONST Int type=kotlin.Int value=0
+      VAR name:l0 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
+            BLOCK_BODY
+              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                BLOCK type=kotlin.Int origin=POSTFIX_INCR
+                  VAR IR_TEMPORARY_VARIABLE name:tmp_0 type:kotlin.Int [val]
+                    GET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Int origin=null
+                  SET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Unit origin=POSTFIX_INCR
+                    CALL 'public final fun inc (): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null
+                      $this: GET_VAR 'val tmp_0: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+                  GET_VAR 'val tmp_0: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+              GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+      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
+            BLOCK_BODY
+              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                BLOCK type=kotlin.Int origin=POSTFIX_INCR
+                  VAR IR_TEMPORARY_VARIABLE name:tmp_1 type:kotlin.Int [val]
+                    GET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Int origin=null
+                  SET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Unit origin=POSTFIX_INCR
+                    CALL 'public final fun inc (): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null
+                      $this: GET_VAR 'val tmp_1: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+                  GET_VAR 'val tmp_1: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+              GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+      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
+            BLOCK_BODY
+              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                BLOCK type=kotlin.Int origin=POSTFIX_INCR
+                  VAR IR_TEMPORARY_VARIABLE name:tmp_2 type:kotlin.Int [val]
+                    GET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Int origin=null
+                  SET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Unit origin=POSTFIX_INCR
+                    CALL 'public final fun inc (): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null
+                      $this: GET_VAR 'val tmp_2: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+                  GET_VAR 'val tmp_2: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+      VAR name:l3 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
+                  VAR IR_TEMPORARY_VARIABLE name:tmp_3 type:kotlin.Int [val]
+                    GET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Int origin=null
+                  SET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Unit origin=POSTFIX_INCR
+                    CALL 'public final fun inc (): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null
+                      $this: GET_VAR 'val tmp_3: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+                  GET_VAR 'val tmp_3: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+              GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+      CALL 'public final fun runUsual (block: kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+        block: GET_VAR 'val l0: kotlin.Function0<kotlin.Unit> declared in <root>.test_1' type=kotlin.Function0<kotlin.Unit> origin=null
+      CALL 'public final fun runUsual (block: 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
+            BLOCK_BODY
+              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                BLOCK type=kotlin.Int origin=POSTFIX_INCR
+                  VAR IR_TEMPORARY_VARIABLE name:tmp_4 type:kotlin.Int [val]
+                    GET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Int origin=null
+                  SET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Unit origin=POSTFIX_INCR
+                    CALL 'public final fun inc (): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null
+                      $this: GET_VAR 'val tmp_4: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+                  GET_VAR 'val tmp_4: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+      CALL 'public final fun runComposable (block: @[MyComposable] kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+        block: GET_VAR 'val l0: kotlin.Function0<kotlin.Unit> declared in <root>.test_1' type=kotlin.Function0<kotlin.Unit> origin=null
+      CALL 'public final fun runComposable (block: @[MyComposable] kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+        block: GET_VAR 'val l1: kotlin.Function0<kotlin.Unit> declared in <root>.test_1' type=kotlin.Function0<kotlin.Unit> origin=null
+      CALL 'public final fun runComposable (block: @[MyComposable] kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+        block: GET_VAR 'val l2: @[MyComposable] kotlin.Function0<kotlin.Unit> declared in <root>.test_1' type=@[MyComposable] kotlin.Function0<kotlin.Unit> origin=null
+      CALL 'public final fun runComposable (block: @[MyComposable] kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+        block: GET_VAR 'val l3: kotlin.Function0<kotlin.Unit> declared in <root>.test_1' type=kotlin.Function0<kotlin.Unit> origin=null
+      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
+            BLOCK_BODY
+              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                BLOCK type=kotlin.Int origin=POSTFIX_INCR
+                  VAR IR_TEMPORARY_VARIABLE name:tmp_5 type:kotlin.Int [val]
+                    GET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Int origin=null
+                  SET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Unit origin=POSTFIX_INCR
+                    CALL 'public final fun inc (): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null
+                      $this: GET_VAR 'val tmp_5: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+                  GET_VAR 'val tmp_5: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+      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
+                  VAR IR_TEMPORARY_VARIABLE name:tmp_6 type:kotlin.Int [val]
+                    GET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Int origin=null
+                  SET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Unit origin=POSTFIX_INCR
+                    CALL 'public final fun inc (): kotlin.Int declared in kotlin.Int' type=kotlin.Int origin=null
+                      $this: GET_VAR 'val tmp_6: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+                  GET_VAR 'val tmp_6: kotlin.Int declared in <root>.test_1.<anonymous>' type=kotlin.Int origin=null
+      RETURN type=kotlin.Nothing from='public final fun test_1 (): kotlin.Int declared in <root>'
+        GET_VAR 'var x: kotlin.Int declared in <root>.test_1' type=kotlin.Int origin=null
+  FUN name:runComposable2 visibility:public modality:FINAL <> (block:kotlin.Function1<kotlin.String, kotlin.String>) returnType:kotlin.String
+    VALUE_PARAMETER name:block index:0 type:kotlin.Function1<kotlin.String, kotlin.String>
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun runComposable2 (block: kotlin.Function1<kotlin.String, kotlin.String>): kotlin.String declared in <root>'
+        CALL 'public abstract fun invoke (p1: P1 of kotlin.Function1): R of kotlin.Function1 declared in kotlin.Function1' type=kotlin.String origin=null
+          $this: GET_VAR 'block: kotlin.Function1<kotlin.String, kotlin.String> declared in <root>.runComposable2' type=kotlin.Function1<kotlin.String, kotlin.String> origin=null
+          p1: CONST String type=kotlin.String value="O"
+  FUN name:test_2 visibility:public modality:FINAL <> () returnType:kotlin.String
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun test_2 (): kotlin.String declared in <root>'
+        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
+              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'
+                  CALL 'public final fun plus (other: kotlin.Any?): kotlin.String declared in kotlin.String' type=kotlin.String origin=PLUS
+                    $this: GET_VAR 'it: kotlin.String declared in <root>.test_2.<anonymous>' type=kotlin.String origin=null
+                    other: CONST String type=kotlin.String value="K"
+  FUN name:box visibility:public modality:FINAL <> () returnType:kotlin.String
+    BLOCK_BODY
+      VAR name:res1 type:kotlin.Int [val]
+        CALL 'public final fun test_1 (): kotlin.Int declared in <root>' type=kotlin.Int origin=null
+      WHEN type=kotlin.Unit origin=IF
+        BRANCH
+          if: CALL 'public final fun not (): kotlin.Boolean declared in kotlin.Boolean' type=kotlin.Boolean origin=EXCLEQ
+            $this: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EXCLEQ
+              arg0: GET_VAR 'val res1: kotlin.Int declared in <root>.box' type=kotlin.Int origin=null
+              arg1: CONST Int type=kotlin.Int value=8
+          then: RETURN type=kotlin.Nothing from='public final fun box (): kotlin.String declared in <root>'
+            STRING_CONCATENATION type=kotlin.String
+              CONST String type=kotlin.String value="Fail 1: "
+              GET_VAR 'val res1: kotlin.Int declared in <root>.box' type=kotlin.Int origin=null
+      VAR name:res2 type:kotlin.String [val]
+        CALL 'public final fun test_2 (): kotlin.String declared in <root>' type=kotlin.String origin=null
+      WHEN type=kotlin.Unit origin=IF
+        BRANCH
+          if: CALL 'public final fun not (): kotlin.Boolean declared in kotlin.Boolean' type=kotlin.Boolean origin=EXCLEQ
+            $this: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EXCLEQ
+              arg0: GET_VAR 'val res2: kotlin.String declared in <root>.box' type=kotlin.String origin=null
+              arg1: CONST String type=kotlin.String value="OK"
+          then: RETURN type=kotlin.Nothing from='public final fun box (): kotlin.String declared in <root>'
+            STRING_CONCATENATION type=kotlin.String
+              CONST String type=kotlin.String value="Fail 2: "
+              GET_VAR 'val res2: kotlin.String declared in <root>.box' type=kotlin.String origin=null
+      RETURN type=kotlin.Nothing from='public final fun box (): kotlin.String declared in <root>'
+        CONST String type=kotlin.String value="OK"

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

@@ -0,0 +1,108 @@
+FILE: composableFunction.kt
+    public final fun runUsual(block: R|() -> kotlin/Unit|): R|kotlin/Unit| {
+        R|<local>/block|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()
+    }
+    public final fun runComposable(block: R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>|): R|kotlin/Unit| {
+        R|<local>/block|.R|SubstitutionOverride<some/MyComposableFunction0.invoke: R|kotlin/Unit|>|()
+    }
+    public final fun test_1(): R|kotlin/Int| {
+        lvar x: R|kotlin/Int| = Int(0)
+        lval l0: R|() -> kotlin/Unit| = fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+             {
+                lval <unary>: R|kotlin/Int| = R|<local>/x|
+                R|<local>/x| = R|<local>/<unary>|.R|kotlin/Int.inc|()
+                R|<local>/<unary>|
+            }
+
+            Q|kotlin/Unit|
+        }
+
+        lval l1: R|some/MyComposableFunction0<kotlin/Unit>| = fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+             {
+                lval <unary>: R|kotlin/Int| = R|<local>/x|
+                R|<local>/x| = R|<local>/<unary>|.R|kotlin/Int.inc|()
+                R|<local>/<unary>|
+            }
+
+            Q|kotlin/Unit|
+        }
+
+        lval l2: R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>| = fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+             {
+                lval <unary>: R|kotlin/Int| = R|<local>/x|
+                R|<local>/x| = R|<local>/<unary>|.R|kotlin/Int.inc|()
+                R|<local>/<unary>|
+            }
+
+        }
+
+        lval l3: R|some/MyComposableFunction0<kotlin/Unit>| = @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+             {
+                lval <unary>: R|kotlin/Int| = R|<local>/x|
+                R|<local>/x| = R|<local>/<unary>|.R|kotlin/Int.inc|()
+                R|<local>/<unary>|
+            }
+
+            Q|kotlin/Unit|
+        }
+
+        R|/runUsual|(R|<local>/l0|)
+        R|/runUsual|(<L> = runUsual@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+             {
+                lval <unary>: R|kotlin/Int| = R|<local>/x|
+                R|<local>/x| = R|<local>/<unary>|.R|kotlin/Int.inc|()
+                R|<local>/<unary>|
+            }
+
+        }
+        )
+        R|/runComposable|(R|<local>/l0|)
+        R|/runComposable|(R|<local>/l1|)
+        R|/runComposable|(R|<local>/l2|)
+        R|/runComposable|(R|<local>/l3|)
+        R|/runComposable|(<L> = runComposable@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+             {
+                lval <unary>: R|kotlin/Int| = R|<local>/x|
+                R|<local>/x| = R|<local>/<unary>|.R|kotlin/Int.inc|()
+                R|<local>/<unary>|
+            }
+
+        }
+        )
+        R|/runComposable|(<L> = @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() runComposable@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+             {
+                lval <unary>: R|kotlin/Int| = R|<local>/x|
+                R|<local>/x| = R|<local>/<unary>|.R|kotlin/Int.inc|()
+                R|<local>/<unary>|
+            }
+
+        }
+        )
+        ^test_1 R|<local>/x|
+    }
+    public final fun runComposable2(block: R|some/MyComposableFunction1<kotlin/String, kotlin/String>|): R|kotlin/String| {
+        ^runComposable2 R|<local>/block|.R|SubstitutionOverride<some/MyComposableFunction1.invoke: R|kotlin/String|>|(String(O))
+    }
+    public final fun test_2(): R|kotlin/String| {
+        ^test_2 R|/runComposable2|(<L> = runComposable2@fun <anonymous>(it: R|kotlin/String|): R|kotlin/String| <inline=NoInline>  {
+            ^ R|<local>/it|.R|kotlin/String.plus|(String(K))
+        }
+        )
+    }
+    public final fun box(): R|kotlin/String| {
+        lval res1: R|kotlin/Int| = R|/test_1|()
+        when () {
+            !=(R|<local>/res1|, Int(8)) ->  {
+                ^box <strcat>(String(Fail 1: ), R|<local>/res1|)
+            }
+        }
+
+        lval res2: R|kotlin/String| = R|/test_2|()
+        when () {
+            !=(R|<local>/res2|, String(OK)) ->  {
+                ^box <strcat>(String(Fail 2: ), R|<local>/res2|)
+            }
+        }
+
+        ^box String(OK)
+    }

plugins/fir-plugin-prototype/testData/box/composableFunction.kt

@@ -0,0 +1,41 @@
+// DUMP_IR
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+
+fun runUsual(block: () -> Unit) { block.invoke() }
+fun runComposable(block: @MyComposable () -> Unit) { block.invoke() }
+
+fun test_1(): Int {
+    var x = 0
+    val l0 = { x++; Unit }
+    val l1: some.MyComposableFunction0<Unit> = { x++; Unit }
+    val l2: @MyComposable (() -> Unit) = { x++ }
+    val l3 = @MyComposable { x++; Unit }
+
+    runUsual(l0)
+    runUsual { x++ }
+
+    runComposable(l0)
+    runComposable(l1)
+    runComposable(l2)
+    runComposable(l3)
+    runComposable { x++ }
+    runComposable @MyComposable { x++ }
+    return x
+}
+
+fun runComposable2(block: some.MyComposableFunction1<String, String>): String {
+    return block.invoke("O")
+}
+
+fun test_2(): String {
+    return runComposable2 { it + "K" }
+}
+
+fun box(): String {
+    val res1 = test_1()
+    if (res1 != 8) return "Fail 1: $res1"
+    val res2 = test_2()
+    if (res2 != "OK") return "Fail 2: $res2"
+    return "OK"
+}

plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirLightTreePluginBlackBoxCodegenTestGenerated.java

@@ -43,6 +43,12 @@ public class FirLightTreePluginBlackBoxCodegenTestGenerated extends AbstractFirL
         runTest("plugins/fir-plugin-prototype/testData/box/classWithGeneratedMembersAndNestedClass.kt");
     }
 
+    @Test
+    @TestMetadata("composableFunction.kt")
+    public void testComposableFunction() throws Exception {
+        runTest("plugins/fir-plugin-prototype/testData/box/composableFunction.kt");
+    }
+
     @Test
     @TestMetadata("generatedClassWithMembersAndNestedClasses.kt")
     public void testGeneratedClassWithMembersAndNestedClasses() throws Exception {