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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBase
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.backend.jvm.ir.parentClassId
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.jvm.config.JvmClasspathRoot
 import org.jetbrains.kotlin.codegen.BytecodeListingTextCollectingVisitor
@@ -27,18 +28,35 @@ 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.IrElement
+import org.jetbrains.kotlin.ir.declarations.IrAnnotationContainer
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrFieldAccessExpression
+import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
 import org.jetbrains.kotlin.ir.util.DumpIrTreeOptions
 import org.jetbrains.kotlin.ir.util.dump
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.FqNameUnsafe
+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 +93,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 annotationToCheckCalls = mainModule.directives[Directives.CHECK_CALLS_WITH_ANNOTATION].singleOrNull()
+        val irCollector = CollectingIrGenerationExtension(annotationToCheckCalls)
 
         val project = mainFile.project
         project.extensionArea.getExtensionPoint(IrGenerationExtension.extensionPointName)
@@ -109,6 +128,12 @@ abstract class AbstractCompilerFacilityTest : AbstractAnalysisApiBasedTest() {
             if (result is KtCompilationResult.Success) {
                 testServices.assertions.assertEqualsToTestDataFileSibling(irCollector.result, extension = ".ir.txt")
             }
+
+            if (annotationToCheckCalls != null) {
+                testServices.assertions.assertEqualsToTestDataFileSibling(
+                    irCollector.functionsWithAnnotationToCheckCalls.joinToString("\n"), extension = ".check_calls.txt"
+                )
+            }
         }
     }
 
@@ -187,6 +212,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_CALLS_WITH_ANNOTATION by stringDirective(
+            "Check whether all functions of calls and getters of properties with a given annotation are listed in *.check_calls.txt or not"
+        )
     }
 }
 
@@ -220,10 +249,12 @@ internal fun createCodeFragment(ktFile: KtFile, module: TestModule, testServices
     }
 }
 
-private class CollectingIrGenerationExtension : IrGenerationExtension {
+private class CollectingIrGenerationExtension(private val annotationToCheckCalls: String?) : IrGenerationExtension {
     lateinit var result: String
         private set
 
+    val functionsWithAnnotationToCheckCalls: MutableSet<String> = mutableSetOf()
+
     override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
         assertFalse { ::result.isInitialized }
 
@@ -235,5 +266,49 @@ private class CollectingIrGenerationExtension : IrGenerationExtension {
         )
 
         result = moduleFragment.dump(dumpOptions)
+
+        annotationToCheckCalls?.let { annotationFqName ->
+            moduleFragment.accept(
+                CheckCallsWithAnnotationVisitor(annotationFqName) { functionsWithAnnotationToCheckCalls.add(it.name.asString()) }, null
+            )
+        }
+    }
+
+    /**
+     * This class recursively visits all calls of functions and getters, and if the function or the getter used for a call has
+     * an annotation whose FqName is [annotationFqName], it runs [handleFunctionWithAnnotation] for the function or the getter.
+     */
+    private class CheckCallsWithAnnotationVisitor(
+        private val annotationFqName: String,
+        private val handleFunctionWithAnnotation: (declaration: IrDeclarationWithName) -> Unit,
+    ) : IrElementVisitorVoid {
+        val annotationClassId by lazy {
+            val annotationFqNameUnsafe = FqNameUnsafe(annotationFqName)
+            ClassId(FqName(annotationFqNameUnsafe.parent()), FqName(annotationFqNameUnsafe.shortName().asString()), false)
+        }
+
+        override fun visitElement(element: IrElement) {
+            element.acceptChildrenVoid(this)
+        }
+
+        override fun visitCall(expression: IrCall) {
+            @OptIn(UnsafeDuringIrConstructionAPI::class)
+            val function = expression.symbol.owner
+            if (function.containsAnnotationToCheckCalls()) {
+                handleFunctionWithAnnotation(function)
+            }
+        }
+
+        override fun visitFieldAccess(expression: IrFieldAccessExpression) {
+            @OptIn(UnsafeDuringIrConstructionAPI::class)
+            val field = expression.symbol.owner
+            if (field.containsAnnotationToCheckCalls()) {
+                handleFunctionWithAnnotation(field)
+            }
+        }
+
+        private fun IrAnnotationContainer.containsAnnotationToCheckCalls() =
+            @OptIn(UnsafeDuringIrConstructionAPI::class)
+            annotations.any { it.symbol.owner.parentClassId == annotationClassId }
     }
 }

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/annotationForFunctionOutOfCodeGenTarget.check_calls.txt

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

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/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/firPluginPrototypeMultiBinaryModule/annotationForFunctionOutOfCodeGenTarget.kt

@@ -0,0 +1,32 @@
+// WITH_FIR_TEST_COMPILER_PLUGIN
+// DUMP_IR
+// CHECK_CALLS_WITH_ANNOTATION: org.jetbrains.kotlin.fir.plugin.MyComposable
+
+// MODULE: main
+// MODULE_KIND: Source
+// 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/firPluginPrototypeMultiBinaryModule/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/firPluginPrototypeMultiBinaryModule/annotationForFunctionOutOfCodeGenTarget2.check_calls.txt

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

analysis/analysis-api/testData/components/compilerFacility/firPluginPrototypeMultiBinaryModule/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.FooKt' type=kotlin.Int origin=GET_PROPERTY
+            CONST String type=kotlin.String value="!"

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

@@ -0,0 +1,28 @@
+// WITH_FIR_TEST_COMPILER_PLUGIN
+// DUMP_IR
+// CHECK_CALLS_WITH_ANNOTATION: org.jetbrains.kotlin.fir.plugin.MyComposable
+
+// MODULE: main
+// MODULE_KIND: Source
+// FILE: main.kt
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+import p3.foo
+
+@MyComposable
+fun Greeting(): String {
+    return "Hi $foo!"
+}
+
+// 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

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.backend
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.backend.generators.*
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.signaturer.FirBasedSignatureComposer
 import org.jetbrains.kotlin.ir.IrLock
@@ -62,6 +63,24 @@ interface Fir2IrComponents {
 
     val annotationsFromPluginRegistrar: Fir2IrIrGeneratedDeclarationsRegistrar
 
+    /**
+     * A set of FIR files serving as input for the fir2ir ([Fir2IrConverter.createIrModuleFragment] function) for conversion to IR.
+     *
+     * We set annotations for IR objects, such as IrFunction, in two scenarios:
+     *  1. For FIR declared in library or precompiled: when creating IR object from FIR
+     *  2. For FIR declared in a source module: when filling contents of IR object in Fir2IrVisitor
+     *
+     * Since Fir2IrVisitor will recursively visit all FIR objects and generate IR objects for them, we handle the first scenario
+     * above as a corner case.
+     *
+     * However, when we use CodeGen analysis API, even FIRs declared in the source module can be out of the compile target files,
+     * because we can run the CodeGen only for a few files of the source module. We use [filesBeingCompiled] for that case
+     * to determine whether a given FIR is declared in a source file to be compiled or not for the CodeGen API. If it is not
+     * declared in a file to be compiled (i.e., target of CodeGen), we have to set annotations for IR when creating its IR like
+     * the first scenario above. We set [filesBeingCompiled] as `null` if we do not use the CodeGen analysis API.
+     */
+    val filesBeingCompiled: Set<FirFile>?
+
     interface Manglers {
         val irMangler: KotlinMangler.IrMangler
         val firMangler: FirMangler

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.backend
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.backend.generators.*
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.descriptors.FirModuleDescriptor
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.signaturer.FirBasedSignatureComposer
@@ -25,6 +26,7 @@ class Fir2IrComponentsStorage(
     override val extensions: Fir2IrExtensions,
     override val configuration: Fir2IrConfiguration,
     override val visibilityConverter: Fir2IrVisibilityConverter,
+    override val filesBeingCompiled: Set<FirFile>?,
     irFakeOverrideBuilderProvider: (IrBuiltIns) -> IrFakeOverrideBuilder,
     moduleDescriptor: FirModuleDescriptor,
     commonMemberStorage: Fir2IrCommonMemberStorage,

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

@@ -59,6 +59,7 @@ import org.jetbrains.kotlin.ir.types.IrTypeSystemContext
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class Fir2IrConverter(
     private val moduleDescriptor: FirModuleDescriptor,
@@ -728,15 +729,25 @@ class Fir2IrConverter(
             session.lazyDeclarationResolver.disableLazyResolveContractChecks()
             val moduleDescriptor = FirModuleDescriptor.createSourceModuleDescriptor(session, kotlinBuiltIns)
             val components = Fir2IrComponentsStorage(
-                session, scopeSession, irFactory, fir2IrExtensions, fir2IrConfiguration, visibilityConverter,
+                session,
+                scopeSession,
+                irFactory,
+                fir2IrExtensions,
+                fir2IrConfiguration,
+                visibilityConverter,
+                runIf(fir2IrConfiguration.allowNonCachedDeclarations) { firFiles.toSet() },
                 { irBuiltins ->
                     IrFakeOverrideBuilder(
                         typeContextProvider(irBuiltins),
                         Fir2IrFakeOverrideStrategy(friendModulesMap(session), commonMemberStorage.symbolTable, irMangler),
                         fir2IrExtensions.externalOverridabilityConditions
                     )
                 },
-                moduleDescriptor, commonMemberStorage, irMangler, specialSymbolProvider, initializedIrBuiltIns
+                moduleDescriptor,
+                commonMemberStorage,
+                irMangler,
+                specialSymbolProvider,
+                initializedIrBuiltIns
             )
 
             fir2IrExtensions.registerDeclarations(commonMemberStorage.symbolTable)

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

@@ -26,6 +26,8 @@ import org.jetbrains.kotlin.fir.lazy.Fir2IrLazyProperty
 import org.jetbrains.kotlin.fir.references.toResolvedBaseSymbol
 import org.jetbrains.kotlin.fir.resolve.calls.FirSimpleSyntheticPropertySymbol
 import org.jetbrains.kotlin.fir.resolve.isKFunctionInvoke
+import org.jetbrains.kotlin.fir.resolve.providers.firProvider
+import org.jetbrains.kotlin.fir.resolve.providers.getContainingFile
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
@@ -965,11 +967,19 @@ 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.isDeclaredInFilesBeingCompiled()
         ) {
             annotationGenerator.generate(this, firAnnotationContainer)
         }
     }
 
+    private fun FirAnnotationContainer.isDeclaredInFilesBeingCompiled(): Boolean {
+        if (filesBeingCompiled == null || this !is FirDeclaration) return false
+        return moduleData.session.firProvider.getContainingFile(symbol) !in filesBeingCompiled
+    }
+
     private inline fun <R> convertCatching(element: FirElement, block: () -> R): R {
         try {
             return block()

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

@@ -35,7 +35,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