From f32cdef7604452bf975b55f664acd67c66789058 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Tue, 1 Mar 2022 11:05:57 +0300
Subject: [PATCH] FirCompanionGenerationProcessor: optimize a bit

---
 .../transformers/plugin/FirCompanionGenerationProcessor.kt | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompanionGenerationProcessor.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompanionGenerationProcessor.kt
index cc0a7a0f35b..62582dc1ef7 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompanionGenerationProcessor.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompanionGenerationProcessor.kt
@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.extensions.FirSwitchableExtensionDeclarationsSymbolProvider
 import org.jetbrains.kotlin.fir.extensions.generatedDeclarationsSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
@@ -25,11 +26,15 @@ class FirCompanionGenerationProcessor(
 }
 
 class FirCompanionGenerationTransformer(val session: FirSession) : FirTransformer<Nothing?>() {
+
+    lateinit var generatedDeclarationProvider: FirSwitchableExtensionDeclarationsSymbolProvider
+
     override fun <E : FirElement> transformElement(element: E, data: Nothing?): E {
         return element
     }
 
     override fun transformFile(file: FirFile, data: Nothing?): FirFile {
+        generatedDeclarationProvider = session.generatedDeclarationsSymbolProvider ?: return file
         return file.transformDeclarations(this, data)
     }
 
@@ -40,7 +45,7 @@ class FirCompanionGenerationTransformer(val session: FirSession) : FirTransforme
 
     private fun generateCompanion(regularClass: FirRegularClass) {
         // TODO: add proper error reporting
-        val generatedDeclarationProvider = session.generatedDeclarationsSymbolProvider ?: return
+        generatedDeclarationProvider = session.generatedDeclarationsSymbolProvider ?: return
         val companionClassId = regularClass.classId.createNestedClassId(SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT)
         when (val generatedCompanion = generatedDeclarationProvider.getClassLikeSymbolByClassId(companionClassId)) {
             null -> {}