FIR IDE: Consider COMPANION_GENERATION phase in lazy resolve

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -174,7 +174,7 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
         var currentPhase = FirResolvePhase.IMPORTS
         while (currentPhase < toPhase) {
             currentPhase = currentPhase.next
-            if (currentPhase.pluginPhase) continue
+            if (skipPhaseInLazyResolve(currentPhase)) continue
             if (checkPCE) checkCanceled()
 
             val transformersToApply = designations.mapNotNull {
@@ -358,7 +358,7 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
 
         while (currentPhase < toPhase) {
             currentPhase = currentPhase.next
-            if (currentPhase.pluginPhase) continue
+            if (skipPhaseInLazyResolve(currentPhase)) continue
             if (checkPCE) checkCanceled()
 
             LazyTransformerFactory.createLazyTransformer(
@@ -396,7 +396,7 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
 
         while (currentPhase < FirResolvePhase.BODY_RESOLVE) {
             currentPhase = currentPhase.next
-            if (currentPhase.pluginPhase) continue
+            if (skipPhaseInLazyResolve(currentPhase)) continue
             if (checkPCE) checkCanceled()
 
             LazyTransformerFactory.createLazyTransformer(
@@ -411,6 +411,10 @@ internal class FirLazyDeclarationResolver(private val firFileBuilder: FirFileBui
             ).transformDeclaration(firFileBuilder.firPhaseRunner)
         }
     }
+
+    private fun skipPhaseInLazyResolve(currentPhase: FirResolvePhase): Boolean {
+        return currentPhase == FirResolvePhase.ANNOTATIONS_FOR_PLUGINS
+    }
 }
 
 private fun KtDeclaration.getContainingEnumEntryAsMemberOfEnumEntry(): KtEnumEntry? {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirDesignatedGeneratedCompanionObjectResolveTransformer.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2022 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.low.level.api.fir.transformers
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirPhaseRunner
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDeclarationDesignationWithFile
+import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.ResolveTreeBuilder
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirCompanionGenerationTransformer
+
+internal class LLFirDesignatedGeneratedCompanionObjectResolveTransformer(
+    val designation: FirDeclarationDesignationWithFile,
+    session: FirSession
+) : LLFirLazyTransformer {
+    private val transformer: FirCompanionGenerationTransformer = FirCompanionGenerationTransformer(session)
+
+    override fun transformDeclaration(phaseRunner: LLFirPhaseRunner) {
+        if (designation.declaration.resolvePhase >= FirResolvePhase.COMPANION_GENERATION) return
+
+        ResolveTreeBuilder.resolvePhase(designation.declaration, FirResolvePhase.COMPANION_GENERATION) {
+            phaseRunner.runPhaseWithCustomResolve(FirResolvePhase.COMPANION_GENERATION) {
+                designation.declaration.transform<FirDeclaration, Nothing?>(transformer, null)
+            }
+        }
+
+        LLFirLazyTransformer.updatePhaseDeep(designation.declaration, FirResolvePhase.COMPANION_GENERATION)
+        ensureResolved(designation.declaration)
+        ensureResolvedDeep(designation.declaration)
+    }
+
+    override fun ensureResolved(declaration: FirDeclaration) {
+        check(declaration.resolvePhase >= FirResolvePhase.COMPANION_GENERATION)
+    }
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LazyTransformerFactory.kt

@@ -24,6 +24,10 @@ internal object LazyTransformerFactory {
         firProviderInterceptor: FirProviderInterceptor?,
         checkPCE: Boolean,
     ): LLFirLazyTransformer = when (phase) {
+        FirResolvePhase.COMPANION_GENERATION -> LLFirDesignatedGeneratedCompanionObjectResolveTransformer(
+            designation = designation,
+            session = designation.firFile.moduleData.session,
+        )
         FirResolvePhase.SEALED_CLASS_INHERITORS -> LLFirLazyTransformer.DUMMY
         FirResolvePhase.SUPER_TYPES -> LLFirDesignatedSupertypeResolverTransformer(
             designation = designation,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompanionGenerationProcessor.kt

@@ -24,7 +24,7 @@ class FirCompanionGenerationProcessor(
     override val transformer: FirTransformer<Nothing?> = FirCompanionGenerationTransformer(session)
 }
 
-private class FirCompanionGenerationTransformer(val session: FirSession) : FirTransformer<Nothing?>() {
+class FirCompanionGenerationTransformer(val session: FirSession) : FirTransformer<Nothing?>() {
     override fun <E : FirElement> transformElement(element: E, data: Nothing?): E {
         return element
     }