K2: Use deserialized type annotation for lambda type resolution.

This commit adds code to check whether a deserialized cone type is a
special function type kind or not when resolving the type of a lambda
expression (anonymous function). If it is a special function kind, it
sets the type of lambda based on the special function kind.

^KT-64994 Fixed

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibraryBinaryModuleFirPluginPrototypeMultiBinaryModuleCompilerFacilityTestGenerated.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.AbstractFirPluginPrototypeMultiBinaryModuleCompilerFacilityTest;
+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/firPluginPrototypeMultiBinaryModule")
+@TestDataPath("$PROJECT_ROOT")
+public class FirIdeNormalAnalysisLibraryBinaryModuleFirPluginPrototypeMultiBinaryModuleCompilerFacilityTestGenerated extends AbstractFirPluginPrototypeMultiBinaryModuleCompilerFacilityTest {
+  @NotNull
+  @Override
+  public AnalysisApiTestConfigurator getConfigurator() {
+    return AnalysisApiFirTestConfiguratorFactory.INSTANCE.createConfigurator(
+      new AnalysisApiTestConfiguratorFactoryData(
+        FrontendKind.Fir,
+        TestModuleKind.LibraryBinary,
+        AnalysisSessionMode.Normal,
+        AnalysisApiMode.Ide
+      )
+    );
+  }
+
+  @Test
+  public void testAllFilesPresentInFirPluginPrototypeMultiBinaryModule() {
+    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule"), Pattern.compile("^(.+)\\.(kt)$"), null, true);
+  }
+
+  @Test
+  @TestMetadata("composableFunctionMultiModules.kt")
+  public void testComposableFunctionMultiModules() {
+    runTest("analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/composableFunctionMultiModules.kt");
+  }
+
+  @Test
+  @TestMetadata("composableFunctionMultiModules2.kt")
+  public void testComposableFunctionMultiModules2() {
+    runTest("analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/composableFunctionMultiModules2.kt");
+  }
+}

analysis/analysis-api-impl-base/build.gradle.kts

@@ -24,6 +24,7 @@ dependencies {
     testImplementation(projectTests(":compiler:tests-common"))
     testApi(projectTests(":compiler:test-infrastructure-utils"))
     testApi(projectTests(":compiler:test-infrastructure"))
+    testImplementation(projectTests(":plugins:fir-plugin-prototype"))
     testImplementation(projectTests(":compiler:tests-common-new"))
     testImplementation(projectTests(":analysis:analysis-api-impl-barebone"))
     testImplementation(project(":analysis:symbol-light-classes"))

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/compilerFacility/AbstractCompilerFacilityTest.kt

@@ -26,19 +26,19 @@ import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
+import org.jetbrains.kotlin.fir.plugin.services.PluginRuntimeAnnotationsProvider
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 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.*
-import org.jetbrains.kotlin.test.services.TestServices
-import org.jetbrains.kotlin.test.services.assertions
 import org.jetbrains.org.objectweb.asm.ClassReader
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
@@ -48,6 +48,11 @@ import kotlin.test.assertFalse
 
 abstract class AbstractMultiModuleCompilerFacilityTest : AbstractCompilerFacilityTest()
 
+abstract class AbstractFirPluginPrototypeMultiBinaryModuleCompilerFacilityTest : AbstractCompilerFacilityTest() {
+    override fun extraCustomRuntimeClasspathProviders(): Array<Constructor<RuntimeClasspathProvider>> =
+        arrayOf(::PluginRuntimeAnnotationsProvider)
+}
+
 abstract class AbstractCompilerFacilityTest : AbstractAnalysisApiBasedTest() {
     private companion object {
         private val ALLOWED_ERRORS = listOf(
@@ -107,6 +112,8 @@ abstract class AbstractCompilerFacilityTest : AbstractAnalysisApiBasedTest() {
         }
     }
 
+    open fun extraCustomRuntimeClasspathProviders(): Array<Constructor<RuntimeClasspathProvider>> = emptyArray()
+
     override fun configureTest(builder: TestConfigurationBuilder) {
         super.configureTest(builder)
         with(builder) {
@@ -116,6 +123,7 @@ abstract class AbstractCompilerFacilityTest : AbstractAnalysisApiBasedTest() {
                 +ConfigurationDirectives.WITH_STDLIB
                 +JvmEnvironmentConfigurationDirectives.FULL_JDK
             }
+            useCustomRuntimeClasspathProviders(*extraCustomRuntimeClasspathProviders())
         }
     }
 

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/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] kotlin.Function0<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

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/composableFunctionMultiModules.kt

@@ -0,0 +1,28 @@
+// WITH_FIR_TEST_COMPILER_PLUGIN
+// 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)
+// MODULE_KIND: Source
+// 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
+    )
+}

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/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
+}

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/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] kotlin.Function0<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"

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/composableFunctionMultiModules2.kt

@@ -0,0 +1,34 @@
+// WITH_FIR_TEST_COMPILER_PLUGIN
+// 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)
+// MODULE_KIND: Source
+// 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) {}

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/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
+}

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/TestModuleStructureFactory.kt

@@ -48,7 +48,7 @@ private typealias ModulesByName = Map<String, KtModuleWithFiles>
  *   - For [TestModule] A and B, where A has dependency on B, A will never appears earlier than B in the result list.
  */
 private fun sortInDependencyPostOrder(testModules: List<TestModule>): List<TestModule> {
-    val namesToModules = buildMap { testModules.forEach { put(it.name, it) } }
+    val namesToModules = testModules.associateBy { it.name }
     val notVisited = testModules.toMutableSet()
     val sortedModules = mutableListOf<TestModule>()
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.transformers
 
 import org.jetbrains.kotlin.KtFakeSourceElementKind
+import org.jetbrains.kotlin.builtins.functions.FunctionTypeKind
 import org.jetbrains.kotlin.fakeElement
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
@@ -812,7 +813,10 @@ class FirCallCompletionResultsWriterTransformer(
         }
 
         if (needUpdateLambdaType) {
-            val kind = expectedType?.functionTypeKind(session)
+            // When we get the FunctionTypeKind, we have to check the deserialized ConeType first because it checks both
+            // class ID and annotations. On the other hand, `functionTypeKind()` checks only class ID.
+            val kind = expectedType?.functionTypeKindForDeserializedConeType()
+                ?: expectedType?.functionTypeKind(session)
                 ?: result.typeRef.coneTypeSafe<ConeClassLikeType>()?.functionTypeKind(session)
             result.replaceTypeRef(result.constructFunctionTypeRef(session, kind))
             session.lookupTracker?.recordTypeResolveAsLookup(result.typeRef, result.source, context.file.source)
@@ -822,6 +826,12 @@ class FirCallCompletionResultsWriterTransformer(
         return result
     }
 
+    private fun ConeKotlinType.functionTypeKindForDeserializedConeType(): FunctionTypeKind? {
+        val coneClassLikeType = type as? ConeClassLikeType ?: return null
+        val classId = coneClassLikeType.classId ?: return null
+        return session.functionTypeService.extractSingleExtensionKindForDeserializedConeType(classId, coneClassLikeType.customAnnotations)
+    }
+
     private fun transformImplicitTypeRefInAnonymousFunction(
         anonymousFunction: FirAnonymousFunction,
     ): FirStatement {

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.callRes
 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.AbstractFirPluginPrototypeMultiBinaryModuleCompilerFacilityTest
 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
@@ -116,6 +117,12 @@ internal fun AnalysisApiTestGroup.generateAnalysisApiTests() {
         test(AbstractMultiModuleCompilerFacilityTest::class, filter = testModuleKindIs(TestModuleKind.LibraryBinary)) {
             model("compilationMultiBinaryModule", pattern = TestGeneratorUtil.KT)
         }
+
+        test(
+            AbstractFirPluginPrototypeMultiBinaryModuleCompilerFacilityTest::class, filter = testModuleKindIs(TestModuleKind.LibraryBinary)
+        ) {
+            model("firPluginPrototypeMultiBinaryModule", pattern = TestGeneratorUtil.KT)
+        }
     }
 
     group(filter = testModuleKindIs(TestModuleKind.Source, TestModuleKind.ScriptSource)) {

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"

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 = Foo@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            ^@Foo Unit
+        }
+        )
+    }
+    public final fun box(): R|kotlin/String| {
+        ^box String(OK)
+    }

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

@@ -0,0 +1,28 @@
+// 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
+    )
+}
+
+fun box(): String = "OK"

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: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: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"
+  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' 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"

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> = 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)
+    }

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

@@ -0,0 +1,34 @@
+// 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) {}
+
+fun box(): String = "OK"

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

@@ -61,6 +61,18 @@ public class FirLightTreePluginBlackBoxCodegenTestGenerated extends AbstractFirL
     runTest("plugins/fir-plugin-prototype/testData/box/composableFunction.kt");
   }
 
+  @Test
+  @TestMetadata("composableFunctionMultiModules.kt")
+  public void testComposableFunctionMultiModules() {
+    runTest("plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.kt");
+  }
+
+  @Test
+  @TestMetadata("composableFunctionMultiModules2.kt")
+  public void testComposableFunctionMultiModules2() {
+    runTest("plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules2.kt");
+  }
+
   @Test
   @TestMetadata("expectComposableFunction.kt")
   public void testExpectComposableFunction() {