Cache NameResolver and KlibMetadataClassDataFinder.

^KT-80438

Author: aviroop-123

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/MetadataLibraryBasedSymbolProvider.kt

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.protobuf.GeneratedMessageLite.GeneratedExtension
 import org.jetbrains.kotlin.resolve.KlibCompilerDeserializationConfiguration
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 import org.jetbrains.kotlin.serialization.deserialization.getClassId
+import java.util.HashMap
 
 abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
     session: FirSession,
@@ -53,6 +54,8 @@ abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
     protected val deserializationConfiguration: KlibCompilerDeserializationConfiguration =
         KlibCompilerDeserializationConfiguration(session.languageVersionSettings)
     private val cachedFragments: MutableMap<L, MutableMap<Pair<String, String>, ProtoBuf.PackageFragment>> = mutableMapOf()
+    private val fragmentToNameResolver = HashMap<Int, NameResolver>()
+    private val fragmentToKlibMetadataClassDataFinder = HashMap<Int, KlibMetadataClassDataFinder>()
 
     private fun getPackageFragment(
         resolvedLibrary: L, packageStringName: String, packageMetadataPart: String
@@ -64,6 +67,22 @@ abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
         }
     }
 
+    private fun getNameResolver(fragment: ProtoBuf.PackageFragment): NameResolver {
+        return fragmentToNameResolver.getOrPut(fragment.hashCode()) {
+            NameResolverImpl(
+                fragment.strings,
+                fragment.qualifiedNames,
+            )
+        }
+    }
+
+    private fun getFinder(fragment: ProtoBuf.PackageFragment, resolver: NameResolver): KlibMetadataClassDataFinder {
+        return fragmentToKlibMetadataClassDataFinder.getOrPut(fragment.hashCode()) {
+            // Assumes the fact that the nameResolver depends only on the fragment.
+            KlibMetadataClassDataFinder(fragment, resolver)
+        }
+    }
+
     override fun computePackagePartsInfos(packageFqName: FqName): List<PackagePartsCacheData> {
         val packageStringName = if (packageFqName.isRoot) "" else packageFqName.asString()
 
@@ -78,10 +97,7 @@ abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
 
                 val packageProto = fragment.`package`
 
-                val nameResolver = NameResolverImpl(
-                    fragment.strings,
-                    fragment.qualifiedNames,
-                )
+                val nameResolver = getNameResolver(fragment)
 
                 PackagePartsCacheData(
                     packageProto,
@@ -112,7 +128,7 @@ abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
     @OptIn(SymbolInternals::class)
     override fun extractClassMetadata(classId: ClassId, parentContext: FirDeserializationContext?): ClassMetadataFindResult? {
         forEachFragmentInPackage(classId.packageFqName) { resolvedLibrary, fragment, nameResolver ->
-            val finder = KlibMetadataClassDataFinder(fragment, nameResolver)
+            val finder = getFinder(fragment, nameResolver)
             val classProto = finder.findClassData(classId)?.classProto ?: return@forEachFragmentInPackage
 
             val moduleData = moduleData(resolvedLibrary) ?: return null
@@ -167,10 +183,7 @@ abstract class MetadataLibraryBasedSymbolProvider<L : MetadataLibrary>(
 
                 val fragment = getPackageFragment(resolvedLibrary, packageStringName, packageMetadataPart)
 
-                val nameResolver = NameResolverImpl(
-                    fragment.strings,
-                    fragment.qualifiedNames,
-                )
+                val nameResolver = getNameResolver(fragment)
 
                 f(resolvedLibrary, fragment, nameResolver)
             }

compiler/util-klib-metadata/src/org/jetbrains/kotlin/library/metadata/KlibMetadataClassDataFinder.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.serialization.deserialization.ClassData
 import org.jetbrains.kotlin.serialization.deserialization.ClassDataFinder
 import org.jetbrains.kotlin.serialization.deserialization.getClassId
+import java.util.HashMap
 
 class KlibMetadataClassDataFinder(
     private val fragment: PackageFragment,
@@ -22,12 +23,13 @@ class KlibMetadataClassDataFinder(
 ) : ClassDataFinder {
     val nameList = fragment.getExtension(KlibMetadataProtoBuf.className).orEmpty()
 
+    val classIdToIndex: Map<ClassId, Int> = HashMap<ClassId, Int>().apply {
+        putAll(nameList.mapIndexed { index, value -> nameResolver.getClassId(value) to index }.toMap())
+    }
+
     override fun findClassData(classId: ClassId): ClassData? {
 
-        val index = nameList.indexOfFirst { nameResolver.getClassId(it) == classId }
-        if (index == -1) {
-            return null
-        }
+        val index = classIdToIndex[classId] ?: return null
 
         val foundClass = fragment.getClass_(index) ?: error("Could not find data for serialized class $classId")