From 4b4c5c4e6bb44ec53c893ebdb15a8ae3c7ea6d7d Mon Sep 17 00:00:00 2001
From: Dmitriy Dolovov <Dmitriy.Dolovov@jetbrains.com>
Date: Thu, 22 Feb 2024 19:30:01 +0100
Subject: [PATCH] [FIR] Serialization to KLIB: Clean-up in
 firKlibSerialization.kt

---
 .../fir/serialization/firKlibSerialization.kt | 57 ++++++++++---------
 1 file changed, 30 insertions(+), 27 deletions(-)

diff --git a/compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/firKlibSerialization.kt b/compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/firKlibSerialization.kt
index 82a17d88b4d..ec9dcb75a3e 100644
--- a/compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/firKlibSerialization.kt
+++ b/compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/firKlibSerialization.kt
@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.serialization.SerializableStringTable
 
+// TODO: handle incremental/monolothic (see klib serializer) - maybe externally
 fun serializeSingleFirFile(
     file: FirFile, session: FirSession, scopeSession: ScopeSession,
     actualizedExpectDeclarations: Set<FirDeclaration>?,
@@ -25,53 +26,55 @@ fun serializeSingleFirFile(
     produceHeaderKlib: Boolean = false,
 ): ProtoBuf.PackageFragment {
     val approximator = TypeApproximatorForMetadataSerializer(session)
+
     val packageSerializer = FirElementSerializer.createTopLevel(
         session, scopeSession, serializerExtension,
         approximator,
         languageVersionSettings,
         produceHeaderKlib
     )
-
-    // TODO: handle incremental/monolothic (see klib serializer) - maybe externally
-
-    val packageProto = packageSerializer.packagePartProto(file.packageFqName, listOf(file), actualizedExpectDeclarations).build()
+    val packageProto = packageSerializer.packagePartProto(
+        file.packageFqName,
+        listOf(file),
+        actualizedExpectDeclarations
+    ).build()
 
     val classesProto = mutableListOf<Pair<ProtoBuf.Class, Int>>()
 
-    fun List<FirClassSymbol<*>>.makeClassesProtoWithNested() {
-        val classSymbols = this
-            .filter {
-                it.fir.isNotExpectOrShouldBeSerialized(actualizedExpectDeclarations) &&
-                        it.fir.isNotPrivateOrShouldBeSerialized(produceHeaderKlib)
-            }
-        for (symbol in classSymbols) {
-            val klass = symbol.fir
-            val classSerializer = FirElementSerializer.create(
-                session, scopeSession, klass, serializerExtension, null,
-                approximator, languageVersionSettings, produceHeaderKlib
-            )
-            val index = classSerializer.stringTable.getFqNameIndex(klass)
+    fun FirClass.makeClassProtoWithNested() {
+        if (!isNotExpectOrShouldBeSerialized(actualizedExpectDeclarations) ||
+            !isNotPrivateOrShouldBeSerialized(produceHeaderKlib)
+        ) {
+            return
+        }
 
-            classesProto += classSerializer.classProto(klass).build() to index
-            classSerializer.computeNestedClassifiersForClass(symbol).filterIsInstance<FirClassSymbol<*>>().makeClassesProtoWithNested()
+        val classSerializer = FirElementSerializer.create(
+            session, scopeSession, klass = this, serializerExtension, parentSerializer = null,
+            approximator, languageVersionSettings, produceHeaderKlib
+        )
+        val index = classSerializer.stringTable.getFqNameIndex(this)
+
+        classesProto += classSerializer.classProto(this).build() to index
+
+        for (nestedClassifierSymbol in classSerializer.computeNestedClassifiersForClass(symbol)) {
+            (nestedClassifierSymbol as? FirClassSymbol<*>)?.fir?.makeClassProtoWithNested()
         }
     }
 
     serializerExtension.processFile(file) {
-        file.declarations.mapNotNull { it.symbol as? FirClassSymbol<*> }.makeClassesProtoWithNested()
-    }
-
-    val hasTopLevelDeclarations = file.declarations.any {
-        it is FirMemberDeclaration && it.isNotExpectOrShouldBeSerialized(actualizedExpectDeclarations) &&
-                it.isNotPrivateOrShouldBeSerialized(produceHeaderKlib) &&
-                (it is FirProperty || it is FirSimpleFunction || it is FirTypeAlias)
+        for (declaration in file.declarations) {
+            (declaration as? FirClass)?.makeClassProtoWithNested()
+        }
     }
 
     return buildKlibPackageFragment(
         packageProto,
         classesProto,
         file.packageFqName,
-        hasTopLevelDeclarations && classesProto.isEmpty(),
+        isEmpty = packageProto.functionList.isEmpty() &&
+                packageProto.propertyList.isEmpty() &&
+                packageProto.typeAliasList.isEmpty() &&
+                classesProto.isEmpty(),
         serializerExtension.stringTable as SerializableStringTable
     )
 }