Propagate annotations from inferred function types when serializing

The inferred type lost its annotations and was converted to non-custom function type after serialization. Fixes: [345261077](https://issuetracker.google.com/issues/345261077)
JetBrains · Jun 10, 2024 · 770fe8d · 770fe8d
1 parent 2742b94
commit 770fe8d
Show file tree

Hide file tree

Showing 5 changed files with 89 additions and 6 deletions.
diff --git a/compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/FunctionalTypeUtils.kt b/compiler/fir/providers/src/org/jetbrains/kotlin/fir/types/FunctionalTypeUtils.kt
@@ -8,22 +8,25 @@ package org.jetbrains.kotlin.fir.types
 import org.jetbrains.kotlin.builtins.functions.FunctionTypeKind
 import org.jetbrains.kotlin.builtins.functions.isBasicFunctionOrKFunction
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.computeTypeAttributes
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.hasAnnotation
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
+import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyAnnotationArgumentMapping
 import org.jetbrains.kotlin.fir.originalForSubstitutionOverride
 import org.jetbrains.kotlin.fir.originalOrSelf
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
-import org.jetbrains.kotlin.fir.resolve.scope
-import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.scopes.CallableCopyTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.scopes.processOverriddenFunctions
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.utils.exceptions.withConeTypeEntry
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.types.AbstractTypeChecker
@@ -115,12 +118,26 @@ fun ConeKotlinType.customFunctionTypeToSimpleFunctionType(session: FirSession):
     } else {
         FunctionTypeKind.Function
     }
-    return createFunctionTypeWithNewKind(session, newKind)
+    return createFunctionTypeWithNewKind(
+        session = session,
+        kind = newKind,
+        additionalAnnotations = kind.annotationOnInvokeClassId
+            ?.takeUnless { classId -> type.attributes.customAnnotations.hasAnnotation(classId, session) }
+            ?.let { annotationClassId ->
+                listOf(buildAnnotation {
+                    argumentMapping = FirEmptyAnnotationArgumentMapping
+                    annotationTypeRef = buildResolvedTypeRef {
+                        type = annotationClassId.defaultType(emptyList())
+                    }
+                })
+            } ?: emptyList()
+    )
 }
 
 fun ConeKotlinType.createFunctionTypeWithNewKind(
     session: FirSession,
     kind: FunctionTypeKind,
+    additionalAnnotations: List<FirAnnotation> = emptyList(),
     updateTypeArguments: (Array<out ConeTypeProjection>.() -> Array<out ConeTypeProjection>)? = null,
 ): ConeClassLikeType {
     val expandedType = fullyExpandedType(session)
@@ -129,7 +146,7 @@ fun ConeKotlinType.createFunctionTypeWithNewKind(
     return functionTypeId.toLookupTag().constructClassType(
         updateTypeArguments?.let { typeArguments.updateTypeArguments() } ?: typeArguments,
         isNullable = expandedType.isNullable,
-        attributes = expandedType.attributes
+        attributes = expandedType.attributes.add(additionalAnnotations.computeTypeAttributes(session, shouldExpandTypeAliases = false)),
     )
 }
 

diff --git a/...ts/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposeCrossModuleTests.kt b/...ts/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposeCrossModuleTests.kt
@@ -1278,6 +1278,45 @@ class ComposeCrossModuleTests(useFir: Boolean) : AbstractCodegenTest(useFir) {
         )
     }
 
+    // regression test for https://issuetracker.google.com/issues/345261077
+    @Test
+    fun composableInferredReturnType() {
+        compile(
+            mapOf(
+                "Base" to mapOf(
+                    "base/Base.kt" to """
+                    package base
+
+                    import androidx.compose.runtime.Composable
+                    
+                    @Composable
+                    fun brokenMangling() = @Composable {
+                    }
+                    """
+                ),
+                "Main" to mapOf(
+                    "Main.kt" to """
+                    package main
+
+                    import androidx.compose.runtime.Composable
+                    import base.brokenMangling
+
+                    @Composable
+                    fun AppManglingPreview() {
+                        brokenMangling()()
+                    }
+                    """
+                )
+            ),
+            validate = {
+                assertTrue(
+                    "Composable lambda type should be implicitly resolved.",
+                    it.contains("INVOKESTATIC base/BaseKt.brokenMangling (Landroidx/compose/runtime/Composer;I)Lkotlin/jvm/functions/Function2;"),
+                )
+            }
+        )
+    }
+
     private fun compile(
         modules: Map<String, Map<String, String>>,
         dumpClasses: Boolean = false,

diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.ir.txt b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.ir.txt
@@ -5,6 +5,18 @@ FILE fqName:p3 fileName:/foo.kt
       MyComposable
     VALUE_PARAMETER name:text index:0 type:@[MyComposable] kotlin.Function0<kotlin.Unit>
     BLOCK_BODY
+  FUN name:FooReturn visibility:public modality:FINAL <> () returnType:kotlin.Function0<kotlin.Unit>
+    annotations:
+      MyComposable
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun FooReturn (): kotlin.Function0<kotlin.Unit> declared in p3'
+        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 p3.FooReturn'
+                GET_OBJECT 'CLASS OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
 Module: main
 FILE fqName:<root> fileName:/main.kt
   FUN name:Bar visibility:public modality:FINAL <> () returnType:kotlin.Unit
@@ -19,6 +31,8 @@ FILE fqName:<root> fileName:/main.kt
             BLOCK_BODY
               RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.Bar'
                 GET_OBJECT 'CLASS OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+      CALL 'public abstract fun invoke (): R of kotlin.Function0 declared in kotlin.Function0' type=kotlin.Unit origin=null
+        $this: CALL 'public final fun FooReturn (): @[MyComposable] some.MyComposableFunction0<kotlin.Unit> declared in p3' type=@[MyComposable] some.MyComposableFunction0<kotlin.Unit> origin=null
   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>'

diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.txt b/plugins/fir-plugin-prototype/testData/box/composableFunctionMultiModules.fir.txt
@@ -4,13 +4,20 @@ FILE: foo.kt
 
     @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| {
     }
+    @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() public final fun FooReturn(): R|some/MyComposableFunction0<kotlin/Unit>| {
+        ^FooReturn @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+            ^ Unit
+        }
+
+    }
 Module: main
 FILE: main.kt
     @R|org/jetbrains/kotlin/fir/plugin/MyComposable|() public final fun Bar(): R|kotlin/Unit| {
         R|p3/Foo|(Foo@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
             ^@Foo Unit
         }
         )
+        R|p3/FooReturn|().R|SubstitutionOverride<some/MyComposableFunction0.invoke: R|kotlin/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
@@ -12,17 +12,23 @@ public fun Foo(
     text: @MyComposable () -> Unit,
 ) {}
 
+@MyComposable
+public fun FooReturn(
+) = @MyComposable {}
+
 // MODULE: main(lib)
 // FILE: main.kt
 
 import org.jetbrains.kotlin.fir.plugin.MyComposable
 import p3.Foo
+import p3.FooReturn
 
 @MyComposable
 public fun Bar() {
     Foo(
         text = {}, // @Composable invocations can only happen from the context of a @Composable function
     )
+    FooReturn()()
 }
 
 fun box(): String = "OK"