From c6886fc779d75b7669ae7485626d14adb25ffc14 Mon Sep 17 00:00:00 2001
From: aleksandrina-streltsova <alexandrina.streltsova@jetbrains.com>
Date: Mon, 9 Jan 2023 13:27:36 +0200
Subject: [PATCH] [FIR] Use post-compute in typeAliasCache for deserialized
 symbols

^KTIJ-24156 Fixed
---
 .../AbstractFirDeserializedSymbolProvider.kt  | 19 +++++++++++++++----
 .../deserialization/FirMemberDeserializer.kt  |  8 ++++++--
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/AbstractFirDeserializedSymbolProvider.kt b/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/AbstractFirDeserializedSymbolProvider.kt
index f3670f91ecb..6287fd84436 100644
--- a/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/AbstractFirDeserializedSymbolProvider.kt
+++ b/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/AbstractFirDeserializedSymbolProvider.kt
@@ -64,6 +64,8 @@ abstract class LibraryPathFilter {
 
 typealias DeserializedClassPostProcessor = (FirRegularClassSymbol) -> Unit
 
+typealias DeserializedTypeAliasPostProcessor = (FirTypeAliasSymbol) -> Unit
+
 abstract class AbstractFirDeserializedSymbolProvider(
     session: FirSession,
     val moduleDataProvider: ModuleDataProvider,
@@ -74,7 +76,15 @@ abstract class AbstractFirDeserializedSymbolProvider(
     // ------------------------ Caches ------------------------
 
     private val packagePartsCache = session.firCachesFactory.createCache(::tryComputePackagePartInfos)
-    private val typeAliasCache = session.firCachesFactory.createCache(::findAndDeserializeTypeAlias)
+    private val typeAliasCache: FirCache<ClassId, FirTypeAliasSymbol?, FirDeserializationContext?> =
+        session.firCachesFactory.createCacheWithPostCompute(
+            createValue = { classId, _ -> findAndDeserializeTypeAlias(classId) },
+            postCompute = { _, symbol, postProcessor ->
+                if (postProcessor != null && symbol != null) {
+                    postProcessor.invoke(symbol)
+                }
+            }
+        )
     private val classCache: FirCache<ClassId, FirRegularClassSymbol?, FirDeserializationContext?> =
         session.firCachesFactory.createCacheWithPostCompute(
             createValue = { classId, context -> findAndDeserializeClass(classId, context) },
@@ -120,13 +130,14 @@ abstract class AbstractFirDeserializedSymbolProvider(
         return computePackagePartsInfos(packageFqName)
     }
 
-    private fun findAndDeserializeTypeAlias(classId: ClassId): FirTypeAliasSymbol? {
+    private fun findAndDeserializeTypeAlias(classId: ClassId): Pair<FirTypeAliasSymbol?, DeserializedTypeAliasPostProcessor?> {
         return getPackageParts(classId.packageFqName).firstNotNullOfOrNull { part ->
             val ids = part.typeAliasNameIndex[classId.shortClassName]
             if (ids == null || ids.isEmpty()) return@firstNotNullOfOrNull null
             val aliasProto = part.proto.getTypeAlias(ids.single())
-            part.context.memberDeserializer.loadTypeAlias(aliasProto).symbol
-        }
+            val postProcessor: DeserializedTypeAliasPostProcessor = { part.context.memberDeserializer.loadTypeAlias(aliasProto, it) }
+            FirTypeAliasSymbol(classId) to postProcessor
+        } ?: (null to null)
     }
 
     private fun findAndDeserializeClass(
diff --git a/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/FirMemberDeserializer.kt b/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/FirMemberDeserializer.kt
index e9c231f2cdc..288c6e03583 100644
--- a/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/FirMemberDeserializer.kt
+++ b/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/FirMemberDeserializer.kt
@@ -191,11 +191,15 @@ class FirMemberDeserializer(private val c: FirDeserializationContext) {
         return lowSixBits + rest
     }
 
-    fun loadTypeAlias(proto: ProtoBuf.TypeAlias): FirTypeAlias {
+    /**
+     * If loading happens in post-compute, then symbol for type alias is already computed and should be provided in [preComputedSymbol]
+     * @See AbstractFirDeserializedSymbolProvider.findAndDeserializeTypeAlias
+     */
+    fun loadTypeAlias(proto: ProtoBuf.TypeAlias, preComputedSymbol: FirTypeAliasSymbol? = null): FirTypeAlias {
         val flags = proto.flags
         val name = c.nameResolver.getName(proto.name)
         val classId = ClassId(c.packageFqName, name)
-        val symbol = FirTypeAliasSymbol(classId)
+        val symbol = preComputedSymbol ?: FirTypeAliasSymbol(classId)
         val local = c.childContext(proto.typeParameterList, containingDeclarationSymbol = symbol)
         return buildTypeAlias {
             moduleData = c.moduleData