K2 AA CodeGen: Set annotation for declarations out of CodeGen targets

The Fir2Ir has two steps in terms of the annotation setup:
 Step #1. Creating `IrMutableAnnotationContainer`
 Step #2. Filling `IrMutableAnnotationContainer` with annotations

When the `IrMutableAnnotationContainer` is out of the source module,
`Fir2IrCallableDeclarationsGenerator` sets annotations for a
`IrMutableAnnotationContainer` when creating its IR object. For example,
when `foo()` defined in `B.kt` is called in file `A.kt`, and `A.kt` is a
part of the source module while `B.kt` is a part of library,
`Fir2IrCallableDeclarationsGenerator` sets annotations of `IrFunction`
for `foo()`.

On the other hand, if `foo()` is a part of the source module,
`Fir2IrCallableDeclarationsGenerator` does not set annotations for it
when creating `IrMutableAnnotationContainer` (Step #1).  Later, when it
fills contents of `IrMutableAnnotationContainer`, it conducts Step #2.

However, if `foo()` is a part of the source module, but it is not in a
compile target file, filling contents of `foo()` does not happen
(because it is not a compile target). As a result, set up annotations
for `foo()` is not done by Fir2Ir. This missing annotation setup causes
a serious bug for Android Compose app.

This commit updates code to decide whether we have to set up the
annotations of `IrMutableAnnotationContainer` in Step #1 or not. We have
to check not only if it is a part of the source code but also if it is a
part of compile targets or not.

^KT-66532 Fixed

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

@@ -26,19 +26,26 @@ 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.ComposableCallVisitor
 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.psi.KtCodeFragment
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtPsiFactory
 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.*
+import org.jetbrains.kotlin.test.services.EnvironmentConfigurator
+import org.jetbrains.kotlin.test.services.RuntimeClasspathProvider
+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
@@ -75,7 +82,8 @@ abstract class AbstractCompilerFacilityTest : AbstractAnalysisApiBasedTest() {
     override fun doTestByMainFile(mainFile: KtFile, mainModule: TestModule, testServices: TestServices) {
         val testFile = mainModule.files.single { it.name == mainFile.name }
 
-        val irCollector = CollectingIrGenerationExtension()
+        val checkComposableFunctions = mainModule.directives.contains(Directives.CHECK_COMPOSABLE_CALL)
+        val irCollector = CollectingIrGenerationExtension(checkComposableFunctions)
 
         val project = mainFile.project
         project.extensionArea.getExtensionPoint(IrGenerationExtension.extensionPointName)
@@ -109,6 +117,12 @@ abstract class AbstractCompilerFacilityTest : AbstractAnalysisApiBasedTest() {
             if (result is KtCompilationResult.Success) {
                 testServices.assertions.assertEqualsToTestDataFileSibling(irCollector.result, extension = ".ir.txt")
             }
+
+            if (checkComposableFunctions) {
+                testServices.assertions.assertEqualsToTestDataFileSibling(
+                    irCollector.composableFunctions.joinToString("\n"), extension = ".composable.txt"
+                )
+            }
         }
     }
 
@@ -187,6 +201,10 @@ abstract class AbstractCompilerFacilityTest : AbstractAnalysisApiBasedTest() {
         val ATTACH_DUPLICATE_STDLIB by directive(
             "Attach the 'stdlib-jvm-minimal-for-test' library to simulate duplicate stdlib dependency"
         )
+
+        val CHECK_COMPOSABLE_CALL by directive(
+            "Check whether all functions of calls and getters of properties with MyComposable annotation are listed in *.composable.txt or not"
+        )
     }
 }
 
@@ -220,10 +238,12 @@ internal fun createCodeFragment(ktFile: KtFile, module: TestModule, testServices
     }
 }
 
-private class CollectingIrGenerationExtension : IrGenerationExtension {
+private class CollectingIrGenerationExtension(private val collectComposableFunctions: Boolean) : IrGenerationExtension {
     lateinit var result: String
         private set
 
+    val composableFunctions: MutableSet<String> = mutableSetOf()
+
     override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
         assertFalse { ::result.isInitialized }
 
@@ -235,5 +255,9 @@ private class CollectingIrGenerationExtension : IrGenerationExtension {
         )
 
         result = moduleFragment.dump(dumpOptions)
+
+        if (collectComposableFunctions) {
+            moduleFragment.accept(ComposableCallVisitor { composableFunctions.add(it.name.asString()) }, null)
+        }
     }
 }

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/annotationForFunctionOutOfCodeGenTarget.composable.txt

@@ -0,0 +1 @@
+BookmarkButton

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/annotationForFunctionOutOfCodeGenTarget.ir.txt

@@ -0,0 +1,12 @@
+MODULE_FRAGMENT
+  FILE fqName:<root> fileName:main.kt
+    FUN name:PostCardSimple visibility:public modality:FINAL <> (navigateToArticle:kotlin.Function1<kotlin.String, kotlin.Unit>, isFavorite:kotlin.Boolean, onToggleFavorite:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit
+      annotations:
+        MyComposable
+      VALUE_PARAMETER name:navigateToArticle index:0 type:kotlin.Function1<kotlin.String, kotlin.Unit>
+      VALUE_PARAMETER name:isFavorite index:1 type:kotlin.Boolean
+      VALUE_PARAMETER name:onToggleFavorite index:2 type:kotlin.Function0<kotlin.Unit>
+      BLOCK_BODY
+        CALL 'public final fun BookmarkButton (isBookmarked: kotlin.Boolean, onClick: kotlin.Function0<kotlin.Unit>): kotlin.Unit declared in p3.JetnewsIconsKt' type=kotlin.Unit origin=null
+          isBookmarked: GET_VAR 'isFavorite: kotlin.Boolean declared in <root>.PostCardSimple' type=kotlin.Boolean origin=null
+          onClick: GET_VAR 'onToggleFavorite: kotlin.Function0<kotlin.Unit> declared in <root>.PostCardSimple' type=kotlin.Function0<kotlin.Unit> origin=null

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/annotationForFunctionOutOfCodeGenTarget.kt

@@ -0,0 +1,31 @@
+// WITH_FIR_TEST_COMPILER_PLUGIN
+// DUMP_IR
+// CHECK_COMPOSABLE_CALL
+
+// MODULE: main
+// FILE: main.kt
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+import p3.BookmarkButton
+
+@MyComposable
+fun PostCardSimple(
+    navigateToArticle: (String) -> Unit,
+    isFavorite: Boolean,
+    onToggleFavorite: () -> Unit
+) {
+    BookmarkButton(
+        isBookmarked = isFavorite,
+        onClick = onToggleFavorite,
+    )
+}
+// FILE: utils/JetnewsIcons.kt
+package p3
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+
+@MyComposable
+fun BookmarkButton(
+    isBookmarked: Boolean,
+    onClick: () -> Unit,
+) {
+}

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/annotationForFunctionOutOfCodeGenTarget.txt

@@ -0,0 +1,4 @@
+public final class MainKt {
+    // source: 'main.kt'
+    public final static method PostCardSimple(p0: kotlin.jvm.functions.Function1, p1: boolean, p2: kotlin.jvm.functions.Function0): void
+}

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/annotationForFunctionOutOfCodeGenTarget2.composable.txt

@@ -0,0 +1 @@
+<get-foo>

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/annotationForFunctionOutOfCodeGenTarget2.ir.txt

@@ -0,0 +1,11 @@
+MODULE_FRAGMENT
+  FILE fqName:<root> fileName:main.kt
+    FUN name:Greeting visibility:public modality:FINAL <> () returnType:kotlin.String
+      annotations:
+        MyComposable
+      BLOCK_BODY
+        RETURN type=kotlin.Nothing from='public final fun Greeting (): kotlin.String declared in <root>'
+          STRING_CONCATENATION type=kotlin.String
+            CONST String type=kotlin.String value="Hi "
+            CALL 'public final fun <get-foo> (): kotlin.Int declared in p3' type=kotlin.Int origin=GET_PROPERTY
+            CONST String type=kotlin.String value="!"

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/annotationForFunctionOutOfCodeGenTarget2.kt

@@ -0,0 +1,29 @@
+// WITH_FIR_TEST_COMPILER_PLUGIN
+// DUMP_IR
+// CHECK_COMPOSABLE_CALL
+
+// MODULE: lib
+// MODULE_KIND: LibraryBinary
+// FILE: p3/foo.kt
+package p3
+
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+
+private var foo_ = 0
+
+fun setFoo(newFoo: Int) {
+    foo_ = newFoo
+}
+
+val foo: Int
+    @MyComposable get() = foo_ + 1
+
+// MODULE: main(lib)
+// FILE: main.kt
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+import p3.foo
+
+@MyComposable
+fun Greeting(): String {
+    return "Hi $foo!"
+}

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiModule/annotationForFunctionOutOfCodeGenTarget2.txt

@@ -0,0 +1,4 @@
+public final class MainKt {
+    // source: 'main.kt'
+    public final static method Greeting(): java.lang.String
+}

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrComponents.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.ir.linkage.IrProvider
 import org.jetbrains.kotlin.ir.overrides.IrFakeOverrideBuilder
 import org.jetbrains.kotlin.ir.util.KotlinMangler
 import org.jetbrains.kotlin.ir.util.SymbolTable
+import org.jetbrains.kotlin.psi.KtFile
 
 interface Fir2IrComponents {
     val session: FirSession
@@ -62,6 +63,8 @@ interface Fir2IrComponents {
 
     val annotationsFromPluginRegistrar: Fir2IrIrGeneratedDeclarationsRegistrar
 
+    val targetFiles: Set<KtFile>
+
     interface Manglers {
         val irMangler: KotlinMangler.IrMangler
         val firMangler: FirMangler

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrComponentsStorage.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.ir.linkage.IrProvider
 import org.jetbrains.kotlin.ir.overrides.IrFakeOverrideBuilder
 import org.jetbrains.kotlin.ir.util.KotlinMangler
 import org.jetbrains.kotlin.ir.util.SymbolTable
+import org.jetbrains.kotlin.psi.KtFile
 
 class Fir2IrComponentsStorage(
     override val session: FirSession,
@@ -25,6 +26,7 @@ class Fir2IrComponentsStorage(
     override val extensions: Fir2IrExtensions,
     override val configuration: Fir2IrConfiguration,
     override val visibilityConverter: Fir2IrVisibilityConverter,
+    override val targetFiles: Set<KtFile>,
     irFakeOverrideBuilderProvider: (IrBuiltIns) -> IrFakeOverrideBuilder,
     moduleDescriptor: FirModuleDescriptor,
     commonMemberStorage: Fir2IrCommonMemberStorage,

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt

@@ -727,6 +727,7 @@ class Fir2IrConverter(
             val moduleDescriptor = FirModuleDescriptor.createSourceModuleDescriptor(session, kotlinBuiltIns)
             val components = Fir2IrComponentsStorage(
                 session, scopeSession, irFactory, fir2IrExtensions, fir2IrConfiguration, visibilityConverter,
+                firFiles.mapNotNull { it.psi as? KtFile }.toSet(),
                 { irBuiltins ->
                     IrFakeOverrideBuilder(
                         typeContextProvider(irBuiltins),

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/Fir2IrCallableDeclarationsGenerator.kt

@@ -965,6 +965,9 @@ class Fir2IrCallableDeclarationsGenerator(val components: Fir2IrComponents) : Fi
     ) {
         if ((firAnnotationContainer as? FirDeclaration)?.let { it.isFromLibrary || it.isPrecompiled } == true
             || origin == IrDeclarationOrigin.FAKE_OVERRIDE
+            // When `firAnnotationContainer` is not in a compile target file, we will not fill contents for
+            // this annotation container later. Therefore, we have to set its annotations here.
+            || firAnnotationContainer.psi?.let { it.containingFile in targetFiles } != true
         ) {
             annotationGenerator.generate(this, firAnnotationContainer)
         }

plugins/fir-plugin-prototype/plugin-annotations/src/commonMain/kotlin/org/jetbrains/kotlin/fir/plugin/annotations.kt

@@ -38,7 +38,7 @@ annotation class SupertypeWithTypeArgument(val kClass: KClass<*>)
 
 annotation class MetaSupertype
 
-@Target(AnnotationTarget.FUNCTION, AnnotationTarget.TYPE)
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.TYPE, AnnotationTarget.PROPERTY_GETTER)
 annotation class MyComposable
 
 annotation class AllPropertiesConstructor

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

@@ -30,6 +30,8 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
+val MyComposableClassId = ClassId(FqName("org.jetbrains.kotlin.fir.plugin"), FqName("MyComposable"), false)
+
 class ComposableFunctionsTransformer(val pluginContext: IrPluginContext) : IrElementVisitorVoid {
     companion object {
         private val INVOKE = Name.identifier("invoke")
@@ -98,12 +100,10 @@ class ComposableFunctionsTransformer(val pluginContext: IrPluginContext) : IrEle
                     it.packageFqName?.asString() == "some"
         } ?: false
 
-    private val composableClassId = ClassId(FqName("org.jetbrains.kotlin.fir.plugin"), FqName("MyComposable"), false)
-
-    private val composableSymbol = pluginContext.referenceClass(composableClassId)!!
+    private val composableSymbol = pluginContext.referenceClass(MyComposableClassId)!!
 
     private fun IrFunction.mark() {
-        if (!hasAnnotation(composableClassId)) {
+        if (!hasAnnotation(MyComposableClassId)) {
             annotations = annotations + IrConstructorCallImpl.fromSymbolOwner(
                 composableSymbol.owner.defaultType,
                 composableSymbol.constructors.single(),

plugins/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/services/PluginAnnotationsProvider.kt

@@ -5,8 +5,17 @@
 
 package org.jetbrains.kotlin.fir.plugin.services
 
+import org.jetbrains.kotlin.backend.jvm.ir.parentClassId
 import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoot
 import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.IrAnnotationContainer
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrFieldAccessExpression
+import org.jetbrains.kotlin.ir.plugin.MyComposableClassId
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 import org.jetbrains.kotlin.platform.isJs
 import org.jetbrains.kotlin.platform.jvm.isJvm
@@ -45,6 +54,32 @@ class PluginRuntimeAnnotationsProvider(testServices: TestServices) : RuntimeClas
     }
 }
 
+/**
+ * This class recursively visits all calls of functions and getters, and if the function or the getter used for a call has
+ * `MyComposable` annotation, it runs [handleComposable] for the function or the getter.
+ */
+class ComposableCallVisitor(private val handleComposable: (declaration: IrDeclarationWithName) -> Unit) : IrElementVisitorVoid {
+    override fun visitElement(element: IrElement) {
+        element.acceptChildrenVoid(this)
+    }
+
+    override fun visitCall(expression: IrCall) {
+        val function = expression.symbol.owner
+        if (expression.symbol.owner.containsComposableAnnotation()) {
+            handleComposable(function)
+        }
+    }
+
+    override fun visitFieldAccess(expression: IrFieldAccessExpression) {
+        val field = expression.symbol.owner
+        if (expression.symbol.owner.containsComposableAnnotation()) {
+            handleComposable(field)
+        }
+    }
+
+    private fun IrAnnotationContainer.containsComposableAnnotation() =
+        annotations.any { it.symbol.owner.parentClassId == MyComposableClassId }
+}
 
 private const val ANNOTATIONS_JAR_DIR = "plugins/fir-plugin-prototype/plugin-annotations/build/libs/"
 private val JVM_ANNOTATIONS_JAR_FILTER = createFilter("plugin-annotations-jvm", ".jar")