Give an option to prepare IrClass for IrGenerationExtension

Motivation: Compose compiler plugin adds an annotation to specific
classes even when they do not have the annotation in their source code.
If this is done by a new codegen API KtFirCompilerFacility, the added
annotation causes FIR resolution issues. It happens because the
IrGenerationExtension in the Compose compiler plugin uses the annotation
before FIR resolutation and FIR2IR conversion.

This commit introduces `prepareIrClassUsedInsideExtension` function in
IrGenerationExtension interface to allow its child extension to prepare
IrClassSymbol before using it inside the extension. The
`prepareIrClassUsedInsideExtension` function is called by
`KtFirCompilerFacility.compile()` before `generate()` in the extension
is called.
diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt
index cf318cd..ee4715a 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt
@@ -140,23 +140,11 @@
 
         val irGeneratorExtensions = IrGenerationExtension.getInstances(project)
 
-        val dependencyFir2IrResults = dependencyFiles
-            .map(::getFullyResolvedFirFile)
-            .groupBy { it.llFirSession }
-            .map { (dependencySession, dependencyFiles) ->
-                val dependencyConfiguration = configuration
-                    .copy()
-                    .apply {
-                        put(CommonConfigurationKeys.USE_FIR, true)
-                        put(CommonConfigurationKeys.LANGUAGE_VERSION_SETTINGS, dependencySession.languageVersionSettings)
-                    }
+        // Run codegen for dependencies.
+        val dependencyFir2IrResults = dependencyFiles.runFir2Ir(configuration, jvmIrDeserializer, diagnosticReporter, irGeneratorExtensions)
 
-                val dependencyFir2IrExtensions = JvmFir2IrExtensions(dependencyConfiguration, jvmIrDeserializer, JvmIrMangler)
-                runFir2Ir(
-                    dependencySession, dependencyFiles, dependencyFir2IrExtensions,
-                    diagnosticReporter, dependencyConfiguration, irGeneratorExtensions
-                )
-            }
+        // Prepare IrClass used inside each extension.
+        irGeneratorExtensions.forEach { it.prepareIrClass(configuration, jvmIrDeserializer, diagnosticReporter, irGeneratorExtensions) }
 
         val targetConfiguration = configuration
             .copy()
@@ -253,6 +241,45 @@
         }
     }
 
+    private fun List<KtFile>.runFir2Ir(
+        configuration: CompilerConfiguration,
+        irDeserializer: JvmIrDeserializer,
+        diagnosticReporter: DiagnosticReporter,
+        irGeneratorExtensions: List<IrGenerationExtension>,
+    ): List<Fir2IrActualizedResult> =
+        map(::getFullyResolvedFirFile).groupBy { it.llFirSession }.map { (dependencySession, dependencyFiles) ->
+            val dependencyConfiguration = configuration.copy().apply {
+                put(CommonConfigurationKeys.USE_FIR, true)
+                put(CommonConfigurationKeys.LANGUAGE_VERSION_SETTINGS, dependencySession.languageVersionSettings)
+            }
+
+            val dependencyFir2IrExtensions = JvmFir2IrExtensions(dependencyConfiguration, irDeserializer, JvmIrMangler)
+            runFir2Ir(
+                dependencySession,
+                dependencyFiles,
+                dependencyFir2IrExtensions,
+                diagnosticReporter,
+                dependencyConfiguration,
+                irGeneratorExtensions
+            )
+        }
+
+    private fun IrGenerationExtension.prepareIrClass(
+        configuration: CompilerConfiguration,
+        irDeserializer: JvmIrDeserializer,
+        diagnosticReporter: DiagnosticReporter,
+        irGeneratorExtensions: List<IrGenerationExtension>,
+    ) {
+        prepareIrClassUsedInsideExtension { classId ->
+            val firClassSymbol =
+                analysisSession.firSymbolProvider.getClassLikeSymbolByClassId(classId) ?: return@prepareIrClassUsedInsideExtension null
+            val ktFile = firClassSymbol.fir.psi?.containingFile as? KtFile ?: return@prepareIrClassUsedInsideExtension null
+            val result = listOf(ktFile).runFir2Ir(configuration, irDeserializer, diagnosticReporter, irGeneratorExtensions.minus(this))
+                .singleOrNull() ?: return@prepareIrClassUsedInsideExtension null
+            result.pluginContext.referenceClass(classId)
+        }
+    }
+
     private fun runFir2Ir(
         session: LLFirSession,
         firFiles: List<FirFile>,
diff --git a/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestForComposeLikeLoweringGenerated.java b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestForComposeLikeLoweringGenerated.java
new file mode 100644
index 0000000..3fe3b95
--- /dev/null
+++ b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestForComposeLikeLoweringGenerated.java
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2023 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.AbstractCompilerFacilityTestForComposeLikeLowering;
+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/composeLikeLowering")
+@TestDataPath("$PROJECT_ROOT")
+public class FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestForComposeLikeLoweringGenerated extends AbstractCompilerFacilityTestForComposeLikeLowering {
+    @NotNull
+    @Override
+    public AnalysisApiTestConfigurator getConfigurator() {
+        return AnalysisApiFirTestConfiguratorFactory.INSTANCE.createConfigurator(
+            new AnalysisApiTestConfiguratorFactoryData(
+                FrontendKind.Fir,
+                TestModuleKind.LibrarySource,
+                AnalysisSessionMode.Normal,
+                AnalysisApiMode.Ide
+            )
+        );
+    }
+
+    @Test
+    public void testAllFilesPresentInComposeLikeLowering() throws Exception {
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+    }
+
+    @Test
+    @TestMetadata("composeStabilityInferred.kt")
+    public void testComposeStabilityInferred() throws Exception {
+        runTest("analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.kt");
+    }
+}
diff --git a/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleCompilerFacilityTestForComposeLikeLoweringGenerated.java b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleCompilerFacilityTestForComposeLikeLoweringGenerated.java
new file mode 100644
index 0000000..b8f4937
--- /dev/null
+++ b/analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleCompilerFacilityTestForComposeLikeLoweringGenerated.java
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2023 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.AbstractCompilerFacilityTestForComposeLikeLowering;
+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/composeLikeLowering")
+@TestDataPath("$PROJECT_ROOT")
+public class FirIdeNormalAnalysisSourceModuleCompilerFacilityTestForComposeLikeLoweringGenerated extends AbstractCompilerFacilityTestForComposeLikeLowering {
+    @NotNull
+    @Override
+    public AnalysisApiTestConfigurator getConfigurator() {
+        return AnalysisApiFirTestConfiguratorFactory.INSTANCE.createConfigurator(
+            new AnalysisApiTestConfiguratorFactoryData(
+                FrontendKind.Fir,
+                TestModuleKind.Source,
+                AnalysisSessionMode.Normal,
+                AnalysisApiMode.Ide
+            )
+        );
+    }
+
+    @Test
+    public void testAllFilesPresentInComposeLikeLowering() throws Exception {
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+    }
+
+    @Test
+    @TestMetadata("composeStabilityInferred.kt")
+    public void testComposeStabilityInferred() throws Exception {
+        runTest("analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.kt");
+    }
+}
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 a648b69..5ba98e6 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
@@ -68,12 +68,14 @@
         ).map { it.name }
     }
 
+    internal open fun createIrCollectorExtension(): CollectingIrGenerationExtension = CollectingIrGenerationExtension()
+
     override fun doTestByModuleStructure(moduleStructure: TestModuleStructure, testServices: TestServices) {
         val mainModule = moduleStructure.modules.first { it.name == "main" }
         val ktFiles = testServices.ktModuleProvider.getModuleFiles(mainModule).filterIsInstance<KtFile>()
         val ktFile = ktFiles.singleOrNull() ?: ktFiles.first { it.name == "main.kt" }
 
-        val irCollector = CollectingIrGenerationExtension()
+        val irCollector = createIrCollectorExtension()
 
         val project = ktFile.project
         project.extensionArea.getExtensionPoint(IrGenerationExtension.extensionPointName)
@@ -215,7 +217,7 @@
     }
 }
 
-private class CollectingIrGenerationExtension : IrGenerationExtension {
+internal open class CollectingIrGenerationExtension : IrGenerationExtension {
     lateinit var result: String
         private set
 
diff --git a/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/compilerFacility/AbstractCompilerFacilityTestForComposeLikeLowering.kt b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/compilerFacility/AbstractCompilerFacilityTestForComposeLikeLowering.kt
new file mode 100644
index 0000000..239fcdd
--- /dev/null
+++ b/analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/compilerFacility/AbstractCompilerFacilityTestForComposeLikeLowering.kt
@@ -0,0 +1,99 @@
+/*
+ * Copyright 2010-2023 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.impl.base.test.cases.components.compilerFacility
+
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.backend.jvm.ir.isInlineClassType
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.expressions.IrConst
+import org.jetbrains.kotlin.ir.expressions.IrConstKind
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
+import org.jetbrains.kotlin.ir.types.defaultType
+import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+
+abstract class AbstractCompilerFacilityTestForComposeLikeLowering : AbstractCompilerFacilityTest() {
+    override fun createIrCollectorExtension(): CollectingIrGenerationExtension = ComposeLikeIrGenerationExtension()
+}
+
+private val stabilityInferredClassId = ClassId(FqName("androidx.compose.runtime.internal"), FqName("StabilityInferred"), false)
+
+private class ComposeLikeIrGenerationExtension : CollectingIrGenerationExtension() {
+    private var stabilityInferredClass: IrClassSymbol? = null
+
+    override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
+        moduleFragment.transformChildrenVoid(ComposeStabilityAnnotationTransformer(pluginContext, stabilityInferredClass))
+        super.generate(moduleFragment, pluginContext)
+    }
+
+    override fun prepareIrClassUsedInsideExtension(classIdToIrClass: (ClassId) -> IrClassSymbol?) {
+        stabilityInferredClass = classIdToIrClass(stabilityInferredClassId)
+    }
+}
+
+private class ComposeStabilityAnnotationTransformer(
+    private val context: IrPluginContext,
+    private val stabilityInferredClass: IrClassSymbol?,
+) : IrElementTransformerVoid() {
+    @OptIn(UnsafeDuringIrConstructionAPI::class)
+    override fun visitClass(declaration: IrClass): IrStatement {
+        val result = super.visitClass(declaration)
+        val cls = result as? IrClass ?: return result
+
+        if (
+            cls.visibility != DescriptorVisibilities.PUBLIC ||
+            cls.isEnumClass ||
+            cls.isEnumEntry ||
+            cls.isInterface ||
+            cls.isAnnotationClass ||
+            cls.isAnonymousObject ||
+            cls.isExpect ||
+            cls.isInner ||
+            cls.isFileClass ||
+            cls.isCompanion ||
+            cls.defaultType.isInlineClassType()
+        ) return cls
+
+        val stabilityInferred = stabilityInferredClass ?: getTopLevelClass(stabilityInferredClassId)
+
+        cls.annotations += IrConstructorCallImpl(
+            UNDEFINED_OFFSET,
+            UNDEFINED_OFFSET,
+            stabilityInferred.defaultType,
+            stabilityInferred.constructors.first(),
+            0,
+            0,
+            1,
+            null
+        ).also {
+            it.putValueArgument(0, irConst(0))
+        }
+        return result
+    }
+
+    private fun irConst(value: Int): IrConst<Int> = IrConstImpl(
+        UNDEFINED_OFFSET,
+        UNDEFINED_OFFSET,
+        context.irBuiltIns.intType,
+        IrConstKind.Int,
+        value
+    )
+
+    private fun getTopLevelClass(classId: ClassId): IrClassSymbol {
+        return context.referenceClass(classId)
+            ?: error("Class not found in the classpath: ${classId.asSingleFqName()}")
+    }
+}
\ No newline at end of file
diff --git a/analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.ir.txt b/analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.ir.txt
new file mode 100644
index 0000000..2b232ced
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.ir.txt
@@ -0,0 +1,23 @@
+MODULE_FRAGMENT
+  FILE fqName:com.test fileName:main.kt
+    CLASS CLASS name:A modality:FINAL visibility:public superTypes:[kotlin.Any]
+      annotations:
+        StabilityInferred(parameters = 0)
+      $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:com.test.A
+      CONSTRUCTOR visibility:public <> () returnType:com.test.A [primary]
+        BLOCK_BODY
+          DELEGATING_CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in kotlin.Any'
+          INSTANCE_INITIALIZER_CALL classDescriptor='CLASS CLASS name:A modality:FINAL visibility:public superTypes:[kotlin.Any]'
+      FUN FAKE_OVERRIDE name:equals visibility:public modality:OPEN <> ($this:kotlin.Any, other:kotlin.Any?) returnType:kotlin.Boolean [fake_override,operator]
+        overridden:
+          public open fun equals (other: kotlin.Any?): kotlin.Boolean [operator] declared in kotlin.Any
+        $this: VALUE_PARAMETER name:<this> type:kotlin.Any
+        VALUE_PARAMETER name:other index:0 type:kotlin.Any?
+      FUN FAKE_OVERRIDE name:hashCode visibility:public modality:OPEN <> ($this:kotlin.Any) returnType:kotlin.Int [fake_override]
+        overridden:
+          public open fun hashCode (): kotlin.Int declared in kotlin.Any
+        $this: VALUE_PARAMETER name:<this> type:kotlin.Any
+      FUN FAKE_OVERRIDE name:toString visibility:public modality:OPEN <> ($this:kotlin.Any) returnType:kotlin.String [fake_override]
+        overridden:
+          public open fun toString (): kotlin.String declared in kotlin.Any
+        $this: VALUE_PARAMETER name:<this> type:kotlin.Any
diff --git a/analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.kt b/analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.kt
new file mode 100644
index 0000000..e9a4501
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.kt
@@ -0,0 +1,11 @@
+// FILE: main.kt
+package com.test
+
+class A() {}
+
+// FILE: lib.kt
+package androidx.compose.runtime.internal
+
+@Target(AnnotationTarget.CLASS)
+@Retention(AnnotationRetention.BINARY)
+annotation class StabilityInferred(val parameters: Int)
\ No newline at end of file
diff --git a/analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.txt b/analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.txt
new file mode 100644
index 0000000..9cd5f04
--- /dev/null
+++ b/analysis/analysis-api/testData/components/compilerFacility/composeLikeLowering/composeStabilityInferred.txt
@@ -0,0 +1,4 @@
+public final class com/test/A {
+    // source: 'main.kt'
+    public method <init>(): void
+}
\ No newline at end of file
diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension.kt
index 4125203..b27a699 100644
--- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension.kt
+++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension.kt
@@ -9,6 +9,8 @@
 import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.linkage.IrDeserializer
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.name.ClassId
 
 interface IrGenerationExtension : IrDeserializer.IrLinkerExtension {
     companion object :
@@ -20,6 +22,21 @@
 
     fun getPlatformIntrinsicExtension(backendContext: BackendContext): IrIntrinsicExtension? = null
 
+    /**
+     * A child of [IrGenerationExtension] can optionally generate an [IrClassSymbol] to use it inside the extension by overriding
+     * this function.
+     *
+     * For example, imagine a compiler plugin that must add a `@com.example.Foo` annotation to specific classes even through the
+     * source code of the classes does not have `com.example.Foo` annotation for the classes. To prepare `@com.example.Foo`, we
+     * can write a child of [IrGenerationExtension] and override this function.
+     *
+     * The IR codegen API will call this function and provide [classIdToIrClass] that generates [IrClassSymbol] for a given [ClassId],
+     * similar to the IR codegen for dependencies. [prepareIrClassUsedInsideExtension] in the child of [IrGenerationExtension] can
+     * just call [classIdToIrClass] with `com.example.Foo`. The child extension can keep the generated [IrClassSymbol] and add it
+     * as an annotation to the classes in [generate].
+     */
+    fun prepareIrClassUsedInsideExtension(classIdToIrClass: (ClassId) -> IrClassSymbol?) {}
+
     // Returns true if this extension should also be applied in the KAPT stub generation mode in Kotlin/JVM. This mode uses light analysis
     // in the compiler frontend to produce an "API-only" class file which is then converted to a .java stub. Because of the light analysis,
     // the resulting IR does not have function bodies and can contain references to error types. If this method returns true, the extension
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 bf79937..befd5a2 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.AbstractCompilerFacilityTestForComposeLikeLowering
 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
@@ -102,6 +103,10 @@
         test(AbstractMultiModuleCompilerFacilityTest::class, filter = testModuleKindIs(TestModuleKind.Source)) {
             model("compilationMultiModule", pattern = TestGeneratorUtil.KT_WITHOUT_DOTS_IN_NAME)
         }
+
+        test(AbstractCompilerFacilityTestForComposeLikeLowering::class) {
+            model("composeLikeLowering", pattern = TestGeneratorUtil.KT_WITHOUT_DOTS_IN_NAME)
+        }
     }
 
     group(filter = testModuleKindIs(TestModuleKind.Source, TestModuleKind.ScriptSource)) {