K2: Attach annotations to FirAnonymousFunction from expected FunctionTypeKind

When we register `MyComposable` as a special FunctionTypeKind, the
following code must set the annotation `MyComposable` for the lambda
expression used for an argument for `content` parameter of `setContent()`:

```
// MODULE: lib
// FILE: p3/foo.kt

package p3;

import org.jetbrains.kotlin.fir.plugin.MyComposable

fun setContent(content: @MyComposable () -> Unit): Int {
    content()
    return 3
}

// MODULE: main(lib)
// FILE: main.kt

import org.jetbrains.kotlin.fir.plugin.MyComposable
import p3.setContent

fun test(): Int {
    return setContent {
        Greeting("test")
    }
}

@MyComposable
fun Greeting(name: String) {
    show("hi $name!")
}
```

The existing KT compiler correctly gets the function type of the lambda
expression used for the argument for `content` parameter i.e.,
`{ Greeting("test") }`, but it does not attach the annotation
`@MyComposable` from the expected function type to the lambda
expression. This commit fixes it.
diff --git a/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleFirPluginPrototypeMultiModuleCompilerFacilityTestGenerated.java b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleFirPluginPrototypeMultiModuleCompilerFacilityTestGenerated.java
new file mode 100644
index 0000000..addf231
--- /dev/null
+++ b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleFirPluginPrototypeMultiModuleCompilerFacilityTestGenerated.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright 2010-2024 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.analysis.api.fir.test.cases.generated.cases.components.compilerFacility;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.util.KtTestUtil;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.analysis.api.fir.test.configurators.AnalysisApiFirTestConfiguratorFactory;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfiguratorFactoryData;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfigurator;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.TestModuleKind;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.FrontendKind;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisSessionMode;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiMode;
+import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.compilerFacility.AbstractFirPluginPrototypeMultiModuleCompilerFacilityTest;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.analysis.api.GenerateAnalysisApiTestsKt}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule")
+@TestDataPath("$PROJECT_ROOT")
+public class FirIdeNormalAnalysisSourceModuleFirPluginPrototypeMultiModuleCompilerFacilityTestGenerated extends AbstractFirPluginPrototypeMultiModuleCompilerFacilityTest {
+    @NotNull
+    @Override
+    public AnalysisApiTestConfigurator getConfigurator() {
+        return AnalysisApiFirTestConfiguratorFactory.INSTANCE.createConfigurator(
+            new AnalysisApiTestConfiguratorFactoryData(
+                FrontendKind.Fir,
+                TestModuleKind.Source,
+                AnalysisSessionMode.Normal,
+                AnalysisApiMode.Ide
+            )
+        );
+    }
+
+    @Test
+    public void testAllFilesPresentInFirPluginPrototypeMultiModule() throws Exception {
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule"), Pattern.compile("^(.+)\\.(kt)$"), null, true);
+    }
+
+    @Test
+    @TestMetadata("composableFunctionMultiModules.kt")
+    public void testComposableFunctionMultiModules() throws Exception {
+        runTest("analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.kt");
+    }
+
+    @Test
+    @TestMetadata("composableFunctionMultiModules2.kt")
+    public void testComposableFunctionMultiModules2() throws Exception {
+        runTest("analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.kt");
+    }
+}
diff --git a/analysis/analysis-api-impl-base/build.gradle.kts b/analysis/analysis-api-impl-base/build.gradle.kts
index 93a34c1..747629f 100644
--- a/analysis/analysis-api-impl-base/build.gradle.kts
+++ b/analysis/analysis-api-impl-base/build.gradle.kts
@@ -22,6 +22,7 @@
     testImplementation(projectTests(":compiler:tests-common"))
     testApi(projectTests(":compiler:test-infrastructure-utils"))
     testApi(projectTests(":compiler:test-infrastructure"))
+    testImplementation(project(":plugins:fir-plugin-prototype"))
     testImplementation(projectTests(":compiler:tests-common-new"))
     testImplementation(projectTests(":analysis:analysis-api-impl-barebone"))
     testImplementation(project(":analysis:symbol-light-classes"))
diff --git a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/compilerFacility/AbstractCompilerFacilityTest.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/compilerFacility/AbstractCompilerFacilityTest.kt
index fb9c190..3478d28 100644
--- a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/compilerFacility/AbstractCompilerFacilityTest.kt
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/compilerFacility/AbstractCompilerFacilityTest.kt
@@ -13,6 +13,9 @@
 import org.jetbrains.kotlin.analysis.api.components.KtCompilerTarget
 import org.jetbrains.kotlin.analysis.api.diagnostics.KtDiagnostic
 import org.jetbrains.kotlin.analysis.api.diagnostics.KtDiagnosticWithPsi
+import org.jetbrains.kotlin.analysis.api.impl.base.test.configurators.ExtensionRegistrarConfigurator
+import org.jetbrains.kotlin.analysis.api.impl.base.test.configurators.PluginAnnotationsProvider
+import org.jetbrains.kotlin.analysis.api.impl.base.test.configurators.PluginRuntimeAnnotationsProvider
 import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedTest
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
@@ -30,15 +33,14 @@
 import org.jetbrains.kotlin.ir.util.DumpIrTreeOptions
 import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.test.Constructor
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.directives.ConfigurationDirectives
 import org.jetbrains.kotlin.test.directives.JvmEnvironmentConfigurationDirectives
 import org.jetbrains.kotlin.test.directives.model.DirectiveApplicability
 import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
 import org.jetbrains.kotlin.test.model.TestModule
-import org.jetbrains.kotlin.test.services.EnvironmentConfigurator
-import org.jetbrains.kotlin.test.services.TestServices
-import org.jetbrains.kotlin.test.services.assertions
+import org.jetbrains.kotlin.test.services.*
 import org.jetbrains.org.objectweb.asm.ClassReader
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
@@ -48,6 +50,14 @@
 
 abstract class AbstractMultiModuleCompilerFacilityTest : AbstractCompilerFacilityTest()
 
+abstract class AbstractFirPluginPrototypeMultiModuleCompilerFacilityTest : AbstractCompilerFacilityTest() {
+    override fun extraConfigurators(): Array<Constructor<AbstractEnvironmentConfigurator>> =
+        arrayOf(::PluginAnnotationsProvider, ::ExtensionRegistrarConfigurator)
+
+    override fun extraCustomRuntimeClasspathProviders(): Array<Constructor<RuntimeClasspathProvider>> =
+        arrayOf(::PluginRuntimeAnnotationsProvider)
+}
+
 abstract class AbstractCompilerFacilityTest : AbstractAnalysisApiBasedTest() {
     private companion object {
         private val ALLOWED_ERRORS = listOf(
@@ -107,15 +117,20 @@
         }
     }
 
+    open fun extraConfigurators(): Array<Constructor<AbstractEnvironmentConfigurator>> = emptyArray()
+
+    open fun extraCustomRuntimeClasspathProviders(): Array<Constructor<RuntimeClasspathProvider>> = emptyArray()
+
     override fun configureTest(builder: TestConfigurationBuilder) {
         super.configureTest(builder)
         with(builder) {
             useDirectives(Directives)
-            useConfigurators(::CompilerFacilityEnvironmentConfigurator)
+            useConfigurators(*(extraConfigurators() + ::CompilerFacilityEnvironmentConfigurator))
             defaultDirectives {
                 +ConfigurationDirectives.WITH_STDLIB
                 +JvmEnvironmentConfigurationDirectives.FULL_JDK
             }
+            useCustomRuntimeClasspathProviders(*extraCustomRuntimeClasspathProviders())
         }
     }
 
diff --git a/plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/services/ExtensionRegistrarConfigurator.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/configurators/ExtensionRegistrarConfigurator.kt
similarity index 86%
rename from plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/services/ExtensionRegistrarConfigurator.kt
rename to analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/configurators/ExtensionRegistrarConfigurator.kt
index 71725a1..c35818b 100644
--- a/plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/services/ExtensionRegistrarConfigurator.kt
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/configurators/ExtensionRegistrarConfigurator.kt
@@ -3,10 +3,11 @@
  * 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.plugin.services
+package org.jetbrains.kotlin.analysis.api.impl.base.test.configurators
 
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar.ExtensionStorage
+import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
 import org.jetbrains.kotlin.fir.plugin.FirPluginPrototypeExtensionRegistrar
@@ -16,6 +17,7 @@
 import org.jetbrains.kotlin.test.services.TestServices
 
 class ExtensionRegistrarConfigurator(testServices: TestServices) : EnvironmentConfigurator(testServices) {
+    @OptIn(ExperimentalCompilerApi::class)
     override fun ExtensionStorage.registerCompilerExtensions(module: TestModule, configuration: CompilerConfiguration) {
         FirExtensionRegistrarAdapter.registerExtension(FirPluginPrototypeExtensionRegistrar())
         IrGenerationExtension.registerExtension(GeneratedDeclarationsIrBodyFiller())
diff --git a/plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/services/PluginAnnotationsProvider.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/configurators/PluginAnnotationsProvider.kt
similarity index 97%
rename from plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/services/PluginAnnotationsProvider.kt
rename to analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/configurators/PluginAnnotationsProvider.kt
index ed4c5f5..7fe849a 100644
--- a/plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/services/PluginAnnotationsProvider.kt
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/configurators/PluginAnnotationsProvider.kt
@@ -3,7 +3,7 @@
  * 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.plugin.services
+package org.jetbrains.kotlin.analysis.api.impl.base.test.configurators
 
 import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoot
 import org.jetbrains.kotlin.config.CompilerConfiguration
diff --git a/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.ir.txt b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.ir.txt
new file mode 100644
index 0000000..bfb0db9
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.ir.txt
@@ -0,0 +1,14 @@
+MODULE_FRAGMENT
+  FILE fqName:<root> fileName:main.kt
+    FUN name:Bar visibility:public modality:FINAL <> () returnType:kotlin.Unit
+      annotations:
+        MyComposable
+      BLOCK_BODY
+        CALL 'public final fun Foo (text: @[MyComposable] some.MyComposableFunction0<kotlin.Unit>): kotlin.Unit declared in p3.FooKt' type=kotlin.Unit origin=null
+          text: 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
+                RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.Bar'
+                  GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
diff --git a/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.kt b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.kt
new file mode 100644
index 0000000..3e46c01
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.kt
@@ -0,0 +1,26 @@
+// DUMP_IR
+
+// MODULE: lib
+// FILE: p3/foo.kt
+
+package p3;
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+
+@MyComposable
+public fun Foo(
+    text: @MyComposable () -> Unit,
+) {}
+
+// MODULE: main(lib)
+// FILE: main.kt
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+import p3.Foo
+
+@MyComposable
+public fun Bar() {
+    Foo(
+        text = {}, // @Composable invocations can only happen from the context of a @Composable function
+    )
+}
\ No newline at end of file
diff --git a/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.txt b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.txt
new file mode 100644
index 0000000..b4464de
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules.txt
@@ -0,0 +1,16 @@
+final class MainKt$Bar$1 {
+    // source: 'main.kt'
+    enclosing method MainKt.Bar()V
+    public final static field INSTANCE: MainKt$Bar$1
+    inner (anonymous) class MainKt$Bar$1
+    static method <clinit>(): void
+    method <init>(): void
+    public synthetic bridge method invoke(): java.lang.Object
+    public final method invoke(): void
+}
+
+public final class MainKt {
+    // source: 'main.kt'
+    inner (anonymous) class MainKt$Bar$1
+    public final static method Bar(): void
+}
diff --git a/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.ir.txt b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.ir.txt
new file mode 100644
index 0000000..58a9f43
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.ir.txt
@@ -0,0 +1,26 @@
+MODULE_FRAGMENT
+  FILE fqName:<root> fileName:main.kt
+    FUN name:Greeting visibility:public modality:FINAL <> (name:kotlin.String) returnType:kotlin.Unit
+      annotations:
+        MyComposable
+      VALUE_PARAMETER name:name index:0 type:kotlin.String
+      BLOCK_BODY
+        CALL 'public final fun show (str: kotlin.String): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+          str: STRING_CONCATENATION type=kotlin.String
+            CONST String type=kotlin.String value="hi "
+            GET_VAR 'name: kotlin.String declared in <root>.Greeting' type=kotlin.String origin=null
+            CONST String type=kotlin.String value="!"
+    FUN name:show visibility:public modality:FINAL <> (str:kotlin.String) returnType:kotlin.Unit
+      VALUE_PARAMETER name:str index:0 type:kotlin.String
+      BLOCK_BODY
+    FUN name:test visibility:public modality:FINAL <> () returnType:kotlin.Int
+      BLOCK_BODY
+        RETURN type=kotlin.Nothing from='public final fun test (): kotlin.Int declared in <root>'
+          CALL 'public final fun setContent (content: @[MyComposable] some.MyComposableFunction0<kotlin.Unit>): kotlin.Int declared in p3.FooKt' type=kotlin.Int origin=null
+            content: 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
+                  CALL 'public final fun Greeting (name: kotlin.String): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                    name: CONST String type=kotlin.String value="test"
diff --git a/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.kt b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.kt
new file mode 100644
index 0000000..3aaa8d8b
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.kt
@@ -0,0 +1,32 @@
+// DUMP_IR
+
+// MODULE: lib
+// FILE: p3/foo.kt
+
+package p3;
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+
+fun setContent(content: @MyComposable () -> Unit): Int {
+    content()
+    return 3
+}
+
+// MODULE: main(lib)
+// FILE: main.kt
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+import p3.setContent
+
+fun test(): Int {
+    return setContent {
+        Greeting("test")
+    }
+}
+
+@MyComposable
+fun Greeting(name: String) {
+    show("hi $name!")
+}
+
+fun show(str: String) {}
\ No newline at end of file
diff --git a/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.txt b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.txt
new file mode 100644
index 0000000..1046687
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/composableFunctionMultiModules2.txt
@@ -0,0 +1,18 @@
+final class MainKt$test$1 {
+    // source: 'main.kt'
+    enclosing method MainKt.test()I
+    public final static field INSTANCE: MainKt$test$1
+    inner (anonymous) class MainKt$test$1
+    static method <clinit>(): void
+    method <init>(): void
+    public synthetic bridge method invoke(): java.lang.Object
+    public final method invoke(): void
+}
+
+public final class MainKt {
+    // source: 'main.kt'
+    inner (anonymous) class MainKt$test$1
+    public final static method Greeting(p0: java.lang.String): void
+    public final static method show(p0: java.lang.String): void
+    public final static method test(): int
+}
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/test/base/compilerPluginConfiguration.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/test/base/compilerPluginConfiguration.kt
index ea4c4d3..c33f84d 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/test/base/compilerPluginConfiguration.kt
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/test/base/compilerPluginConfiguration.kt
@@ -6,13 +6,13 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.test.base
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analysis.api.impl.base.test.configurators.ExtensionRegistrarConfigurator
+import org.jetbrains.kotlin.analysis.api.impl.base.test.configurators.PluginAnnotationsProvider
 import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
 import org.jetbrains.kotlin.config.AnalysisFlag
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.LanguageVersion
-import org.jetbrains.kotlin.fir.plugin.services.ExtensionRegistrarConfigurator
-import org.jetbrains.kotlin.fir.plugin.services.PluginAnnotationsProvider
 import org.jetbrains.kotlin.test.Constructor
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.directives.model.Directive
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt
index d924b1a..f2ce423 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt
@@ -9,6 +9,7 @@
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.resolve.calls.ArgumentTypeMismatch
@@ -61,6 +62,8 @@
             duringCompletion || sink == null
         ) ?: extractLambdaInfo(expectedType, anonymousFunction, csBuilder, context.session, this)
 
+    anonymousFunction.attachMissingAnnotationFromFunctionType(resolvedArgument.expectedFunctionTypeKind, expectedType)
+
     if (expectedType != null) {
         val parameters = resolvedArgument.parameters
         val functionTypeKind = context.session.functionTypeService.extractSingleSpecialKindForFunction(anonymousFunction.symbol)
@@ -94,6 +97,21 @@
     return resolvedArgument
 }
 
+private fun FirAnonymousFunction.attachMissingAnnotationFromFunctionType(
+    expectedFunctionTypeKind: FunctionTypeKind?,
+    expectedType: ConeKotlinType?,
+) {
+    fun FirAnnotation.getClassId() = annotationTypeRef.coneType.classId
+
+    expectedFunctionTypeKind?.annotationOnInvokeClassId?.let { classId ->
+        if (annotations.any { it.getClassId() == classId }) return@let
+        val annotation =
+            expectedType?.attributes?.get(CustomAnnotationTypeAttribute::class)?.annotations?.single { it.getClassId() == classId }
+                ?: return@let
+        replaceAnnotations(annotations + annotation)
+    }
+}
+
 fun Candidate.preprocessCallableReference(
     argument: FirCallableReferenceAccess,
     expectedType: ConeKotlinType?,
diff --git a/generators/analysis-api-generator/tests/org/jetbrains/kotlin/generators/tests/analysis/api/analysisApi.kt b/generators/analysis-api-generator/tests/org/jetbrains/kotlin/generators/tests/analysis/api/analysisApi.kt
index c71084c..af3b340 100644
--- a/generators/analysis-api-generator/tests/org/jetbrains/kotlin/generators/tests/analysis/api/analysisApi.kt
+++ b/generators/analysis-api-generator/tests/org/jetbrains/kotlin/generators/tests/analysis/api/analysisApi.kt
@@ -11,6 +11,7 @@
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.callResolver.AbstractResolveCandidatesTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.compileTimeConstantProvider.AbstractCompileTimeConstantEvaluatorTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.compilerFacility.AbstractCompilerFacilityTest
+import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.compilerFacility.AbstractFirPluginPrototypeMultiModuleCompilerFacilityTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.compilerFacility.AbstractMultiModuleCompilerFacilityTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.containingDeclarationProvider.AbstractContainingDeclarationProviderByDelegatedMemberScopeTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.containingDeclarationProvider.AbstractContainingDeclarationProviderByMemberScopeTest
@@ -112,6 +113,10 @@
         test(AbstractMultiModuleCompilerFacilityTest::class, filter = testModuleKindIs(TestModuleKind.Source)) {
             model("compilationMultiModule", pattern = TestGeneratorUtil.KT)
         }
+
+        test(AbstractFirPluginPrototypeMultiModuleCompilerFacilityTest::class, filter = testModuleKindIs(TestModuleKind.Source)) {
+            model("firPluginPrototypeMultiModule", pattern = TestGeneratorUtil.KT)
+        }
     }
 
     group(filter = testModuleKindIs(TestModuleKind.Source, TestModuleKind.ScriptSource)) {
diff --git a/plugins/fir-plugin-prototype/build.gradle.kts b/plugins/fir-plugin-prototype/build.gradle.kts
index 5414c19..c694a04 100644
--- a/plugins/fir-plugin-prototype/build.gradle.kts
+++ b/plugins/fir-plugin-prototype/build.gradle.kts
@@ -22,6 +22,7 @@
     testApi(platform(libs.junit.bom))
     testImplementation(libs.junit.jupiter.api)
     testRuntimeOnly(libs.junit.jupiter.engine)
+    testApi(projectTests(":analysis:analysis-api-impl-base"))
     testApi(projectTests(":compiler:tests-common-new"))
     testApi(projectTests(":compiler:test-infrastructure"))
     testApi(projectTests(":compiler:test-infrastructure-utils"))
diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.ir.txt b/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.ir.txt
index 959ff25..68c1fc0 100644
--- a/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.ir.txt
+++ b/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.ir.txt
@@ -91,6 +91,8 @@
       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
diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.txt b/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.txt
index 64ac7c0..f61e603 100644
--- a/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.txt
+++ b/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.txt
@@ -60,7 +60,7 @@
         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>  {
+        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|()
diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.ir.txt b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.ir.txt
new file mode 100644
index 0000000..f243628
--- /dev/null
+++ b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.ir.txt
@@ -0,0 +1,25 @@
+Module: lib
+FILE fqName:p3 fileName:/foo.kt
+  FUN name:Foo visibility:public modality:FINAL <> (text:@[MyComposable] kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit
+    annotations:
+      MyComposable
+    VALUE_PARAMETER name:text index:0 type:@[MyComposable] kotlin.Function0<kotlin.Unit>
+    BLOCK_BODY
+Module: main
+FILE fqName:<root> fileName:/main.kt
+  FUN name:Bar visibility:public modality:FINAL <> () returnType:kotlin.Unit
+    annotations:
+      MyComposable
+    BLOCK_BODY
+      CALL 'public final fun Foo (text: @[MyComposable] some.MyComposableFunction0<kotlin.Unit>): kotlin.Unit declared in p3' type=kotlin.Unit origin=null
+        text: 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
+              RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.Bar'
+                GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+  FUN name:box visibility:public modality:FINAL <> () returnType:kotlin.String
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun box (): kotlin.String declared in <root>'
+        CONST String type=kotlin.String value="OK"
diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.txt b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.txt
new file mode 100644
index 0000000..2128253
--- /dev/null
+++ b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.txt
@@ -0,0 +1,17 @@
+Module: lib
+FILE: foo.kt
+    package p3
+
+    @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() public final fun Foo(text: R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>|): R|kotlin/Unit| {
+    }
+Module: main
+FILE: main.kt
+    @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() public final fun Bar(): R|kotlin/Unit| {
+        R|p3/Foo|(text = @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() Foo@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            ^@Foo Unit
+        }
+        )
+    }
+    public final fun box(): R|kotlin/String| {
+        ^box String(OK)
+    }
diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.kt b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.kt
new file mode 100644
index 0000000..b9a26f3
--- /dev/null
+++ b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.kt
@@ -0,0 +1,30 @@
+// DUMP_IR
+// WITH_STDLIB
+// TARGET_BACKEND: JVM_IR
+
+// MODULE: lib
+// FILE: p3/foo.kt
+
+package p3;
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+
+@MyComposable
+public fun Foo(
+    text: @MyComposable () -> Unit,
+) {}
+
+// MODULE: main(lib)
+// FILE: main.kt
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+import p3.Foo
+
+@MyComposable
+public fun Bar() {
+    Foo(
+        text = {}, // @Composable invocations can only happen from the context of a @Composable function
+    )
+}
+
+fun box() = "OK"
\ No newline at end of file
diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.fir.ir.txt b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.fir.ir.txt
new file mode 100644
index 0000000..0d682e2
--- /dev/null
+++ b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.fir.ir.txt
@@ -0,0 +1,39 @@
+Module: lib
+FILE fqName:p3 fileName:/foo.kt
+  FUN name:setContent visibility:public modality:FINAL <> (content:@[MyComposable] kotlin.Function0<kotlin.Unit>) returnType:kotlin.Int
+    VALUE_PARAMETER name:content 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 'content: @[MyComposable] kotlin.Function0<kotlin.Unit> declared in p3.setContent' type=@[MyComposable] kotlin.Function0<kotlin.Unit> origin=VARIABLE_AS_FUNCTION
+      RETURN type=kotlin.Nothing from='public final fun setContent (content: @[MyComposable] kotlin.Function0<kotlin.Unit>): kotlin.Int declared in p3'
+        CONST Int type=kotlin.Int value=3
+Module: main
+FILE fqName:<root> fileName:/main.kt
+  FUN name:test visibility:public modality:FINAL <> () returnType:kotlin.Int
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun test (): kotlin.Int declared in <root>'
+        CALL 'public final fun setContent (content: @[MyComposable] some.MyComposableFunction0<kotlin.Unit>): kotlin.Int declared in p3' type=kotlin.Int origin=null
+          content: 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
+                CALL 'public final fun Greeting (name: kotlin.String): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                  name: CONST String type=kotlin.String value="test"
+  FUN name:Greeting visibility:public modality:FINAL <> (name:kotlin.String) returnType:kotlin.Unit
+    annotations:
+      MyComposable
+    VALUE_PARAMETER name:name index:0 type:kotlin.String
+    BLOCK_BODY
+      CALL 'public final fun show (str: kotlin.String): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+        str: STRING_CONCATENATION type=kotlin.String
+          CONST String type=kotlin.String value="hi "
+          GET_VAR 'name: kotlin.String declared in <root>.Greeting' type=kotlin.String origin=null
+          CONST String type=kotlin.String value="!"
+  FUN name:show visibility:public modality:FINAL <> (str:kotlin.String) returnType:kotlin.Unit
+    VALUE_PARAMETER name:str index:0 type:kotlin.String
+    BLOCK_BODY
+  FUN name:box visibility:public modality:FINAL <> () returnType:kotlin.String
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun box (): kotlin.String declared in <root>'
+        CONST String type=kotlin.String value="OK"
diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.fir.txt b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.fir.txt
new file mode 100644
index 0000000..226d96df
--- /dev/null
+++ b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.fir.txt
@@ -0,0 +1,24 @@
+Module: lib
+FILE: foo.kt
+    package p3
+
+    public final fun setContent(content: R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>|): R|kotlin/Int| {
+        R|<local>/content|.R|SubstitutionOverride<some/MyComposableFunction0.invoke: R|kotlin/Unit|>|()
+        ^setContent Int(3)
+    }
+Module: main
+FILE: main.kt
+    public final fun test(): R|kotlin/Int| {
+        ^test R|p3/setContent|(<L> = @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() setContent@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            R|/Greeting|(String(test))
+        }
+        )
+    }
+    @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() public final fun Greeting(name: R|kotlin/String|): R|kotlin/Unit| {
+        R|/show|(<strcat>(String(hi ), R|<local>/name|, String(!)))
+    }
+    public final fun show(str: R|kotlin/String|): R|kotlin/Unit| {
+    }
+    public final fun box(): R|kotlin/String| {
+        ^box String(OK)
+    }
diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.kt b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.kt
new file mode 100644
index 0000000..622d5f0
--- /dev/null
+++ b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.kt
@@ -0,0 +1,36 @@
+// DUMP_IR
+// WITH_STDLIB
+// TARGET_BACKEND: JVM_IR
+
+// MODULE: lib
+// FILE: p3/foo.kt
+
+package p3;
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+
+fun setContent(content: @MyComposable () -> Unit): Int {
+    content()
+    return 3
+}
+
+// MODULE: main(lib)
+// FILE: main.kt
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+import p3.setContent
+
+fun test(): Int {
+    return setContent {
+        Greeting("test")
+    }
+}
+
+@MyComposable
+fun Greeting(name: String) {
+    show("hi $name!")
+}
+
+fun show(str: String) {}
+
+fun box() = "OK"
\ No newline at end of file
diff --git a/plugins/fir-plugin-prototype/testData/box/expectComposableFunction.fir.txt b/plugins/fir-plugin-prototype/testData/box/expectComposableFunction.fir.txt
index 5cc2658..658e4c5 100644
--- a/plugins/fir-plugin-prototype/testData/box/expectComposableFunction.fir.txt
+++ b/plugins/fir-plugin-prototype/testData/box/expectComposableFunction.fir.txt
@@ -2,7 +2,7 @@
 FILE: module_common_expectComposableFunction.kt
     @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() public final expect fun ExpectComposable(value: R|kotlin/String|, content: R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction1<@R|kotlin/ParameterName|(name = String(v))  kotlin/String, kotlin/String>|): R|kotlin/String|
     public final fun commonBox(): R|kotlin/String| {
-        ^commonBox R|/ExpectComposable|(String(O), <L> = ExpectComposable@fun <anonymous>(it: R|@R|kotlin/ParameterName|(name = String(v))  kotlin/String|): R|kotlin/String| <inline=NoInline>  {
+        ^commonBox R|/ExpectComposable|(String(O), <L> = @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() ExpectComposable@fun <anonymous>(it: R|@R|kotlin/ParameterName|(name = String(v))  kotlin/String|): R|kotlin/String| <inline=NoInline>  {
             ^ R|<local>/it|.R|kotlin/String.plus|(String(K))
         }
         )
@@ -18,7 +18,7 @@
 FILE: module_common_expectComposableFunction.kt
     @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() public final expect fun ExpectComposable(value: R|kotlin/String|, content: R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction1<@R|kotlin/ParameterName|(name = String(v))  kotlin/String, kotlin/String>|): R|kotlin/String|
     public final fun commonBox(): R|kotlin/String| {
-        ^commonBox R|/ExpectComposable|(String(O), <L> = ExpectComposable@fun <anonymous>(it: R|@R|kotlin/ParameterName|(name = String(v))  kotlin/String|): R|kotlin/String| <inline=NoInline>  {
+        ^commonBox R|/ExpectComposable|(String(O), <L> = @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() ExpectComposable@fun <anonymous>(it: R|@R|kotlin/ParameterName|(name = String(v))  kotlin/String|): R|kotlin/String| <inline=NoInline>  {
             ^ R|<local>/it|.R|kotlin/String.plus|(String(K))
         }
         )
diff --git a/plugins/fir-plugin-prototype/testData/box/nullableComposableParameterAcrossModules.fir.txt b/plugins/fir-plugin-prototype/testData/box/nullableComposableParameterAcrossModules.fir.txt
index f471c68..2595a82 100644
--- a/plugins/fir-plugin-prototype/testData/box/nullableComposableParameterAcrossModules.fir.txt
+++ b/plugins/fir-plugin-prototype/testData/box/nullableComposableParameterAcrossModules.fir.txt
@@ -6,7 +6,7 @@
 FILE: b.kt
     public final fun bar(onClick: R|(() -> kotlin/Unit)?| = Null(null)): R|kotlin/Unit| {
         R|/RMenuItem|(onClick = R|<local>/onClick|, trailing = R|<local>/onClick|?.{ $subj$.R|kotlin/let|<R|() -> kotlin/Unit|, R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>?|>(<L> = let@fun <anonymous>(it: R|() -> kotlin/Unit|): R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>?| <inline=Inline, kind=EXACTLY_ONCE>  {
-            ^ let@fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+            ^ @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() let@fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
                 ^@let Unit
             }
 
diff --git a/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.fir.txt b/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.fir.txt
index 8b08f8f..3b4de6c 100644
--- a/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.fir.txt
+++ b/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.fir.txt
@@ -14,7 +14,7 @@
             ^@id Unit
         }
         )
-        R|/composableId|<R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>|>(composableId@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+        R|/composableId|<R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>|>(@R|org/jetbrains/kotlin/fir/plugin/MyComposable|() composableId@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
             ^@composableId Unit
         }
         )
diff --git a/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/simple.fir.txt b/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/simple.fir.txt
index 055f6fb..14a690b 100644
--- a/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/simple.fir.txt
+++ b/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/simple.fir.txt
@@ -36,7 +36,7 @@
         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>  {
+        R|/runComposable|(<L> = @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() runComposable@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
             ^@runComposable Unit
         }
         )
diff --git a/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirLightTreePluginBlackBoxCodegenTestGenerated.java b/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirLightTreePluginBlackBoxCodegenTestGenerated.java
index cf21cce..a4faf1d 100644
--- a/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirLightTreePluginBlackBoxCodegenTestGenerated.java
+++ b/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirLightTreePluginBlackBoxCodegenTestGenerated.java
@@ -62,6 +62,18 @@
     }
 
     @Test
+    @TestMetadata("composableFunctionMultiModules.kt")
+    public void testComposableFunctionMultiModules() throws Exception {
+        runTest("plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.kt");
+    }
+
+    @Test
+    @TestMetadata("composableFunctionMultiModules2.kt")
+    public void testComposableFunctionMultiModules2() throws Exception {
+        runTest("plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.kt");
+    }
+
+    @Test
     @TestMetadata("expectComposableFunction.kt")
     public void testExpectComposableFunction() throws Exception {
         runTest("plugins/fir-plugin-prototype/testData/box/expectComposableFunction.kt");
diff --git a/plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/runners/PluginPrototypeTests.kt b/plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/runners/PluginPrototypeTests.kt
index 9c967a7..e275d81 100644
--- a/plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/runners/PluginPrototypeTests.kt
+++ b/plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/runners/PluginPrototypeTests.kt
@@ -5,9 +5,9 @@
 
 package org.jetbrains.kotlin.fir.plugin.runners
 
-import org.jetbrains.kotlin.fir.plugin.services.ExtensionRegistrarConfigurator
-import org.jetbrains.kotlin.fir.plugin.services.PluginAnnotationsProvider
-import org.jetbrains.kotlin.fir.plugin.services.PluginRuntimeAnnotationsProvider
+import org.jetbrains.kotlin.analysis.api.impl.base.test.configurators.ExtensionRegistrarConfigurator
+import org.jetbrains.kotlin.analysis.api.impl.base.test.configurators.PluginAnnotationsProvider
+import org.jetbrains.kotlin.analysis.api.impl.base.test.configurators.PluginRuntimeAnnotationsProvider
 import org.jetbrains.kotlin.js.test.fir.AbstractFirLoadK2CompiledJsKotlinTest
 import org.jetbrains.kotlin.test.backend.handlers.IrPrettyKotlinDumpHandler
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder