FIR: use FirCachesFactory in BuiltInsPackageFragment.tryGetSyntheticFunctionalInterface

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirBuiltinSymbolProvider.kt

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.caches.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
 import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
@@ -47,11 +47,8 @@ open class FirBuiltinSymbolProvider(
     val moduleData: FirModuleData,
     val kotlinScopeProvider: FirKotlinScopeProvider
 ) : FirSymbolProvider(session) {
-
-    private data class SyntheticFunctionalInterfaceSymbolKey(val kind: FunctionClassKind, val arity: Int)
-
     private val allPackageFragments = loadBuiltIns().groupBy { it.fqName }
-    private val syntheticFunctionalInterfaceSymbols = mutableMapOf<SyntheticFunctionalInterfaceSymbolKey, FirRegularClassSymbol>()
+    private val syntheticFunctionalInterfaceCache = SyntheticFunctionalInterfaceCache(moduleData, kotlinScopeProvider)
 
     private fun loadBuiltIns(): List<BuiltInsPackageFragment> {
         val classLoader = this::class.java.classLoader
@@ -61,7 +58,7 @@ open class FirBuiltinSymbolProvider(
         return packageFqNames.map { fqName ->
             val resourcePath = BuiltInSerializerProtocol.getBuiltInsFilePath(fqName)
             val inputStream = streamProvider(resourcePath) ?: throw IllegalStateException("Resource not found in classpath: $resourcePath")
-            BuiltInsPackageFragment(inputStream, fqName)
+            BuiltInsPackageFragment(inputStream, fqName, moduleData, kotlinScopeProvider)
         }
     }
 
@@ -73,19 +70,134 @@ open class FirBuiltinSymbolProvider(
     override fun getClassLikeSymbolByFqName(classId: ClassId): FirRegularClassSymbol? {
         return allPackageFragments[classId.packageFqName]?.firstNotNullOfOrNull {
             it.getClassLikeSymbolByFqName(classId)
-        } ?: trySyntheticFunctionalInterface(classId)
+        } ?: syntheticFunctionalInterfaceCache.tryGetSyntheticFunctionalInterface(classId)
     }
 
-    private fun trySyntheticFunctionalInterface(classId: ClassId): FirRegularClassSymbol? {
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
+        allPackageFragments[packageFqName]?.flatMapTo(destination) {
+            it.getTopLevelCallableSymbols(name)
+        }
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
+        getTopLevelFunctionSymbolsToByPackageFragments(destination, packageFqName, name)
+    }
+
+    protected fun getTopLevelFunctionSymbolsToByPackageFragments(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
+        allPackageFragments[packageFqName]?.flatMapTo(destination) {
+            it.getTopLevelFunctionSymbols(name)
+        }
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
+    }
+
+    private class BuiltInsPackageFragment(
+        stream: InputStream, val fqName: FqName, val moduleData: FirModuleData,
+        val kotlinScopeProvider: FirKotlinScopeProvider,
+    ) {
+
+        private val binaryVersionAndPackageFragment = BinaryVersionAndPackageFragment.createFromStream(stream)
+
+        val version: BuiltInsBinaryVersion get() = binaryVersionAndPackageFragment.version
+        val packageProto: ProtoBuf.PackageFragment get() = binaryVersionAndPackageFragment.packageFragment
+
+        private val nameResolver = NameResolverImpl(packageProto.strings, packageProto.qualifiedNames)
+
+        val classDataFinder = ProtoBasedClassDataFinder(packageProto, nameResolver, version) { SourceElement.NO_SOURCE }
+
+        private val memberDeserializer by lazy {
+            FirDeserializationContext.createForPackage(
+                fqName, packageProto.`package`, nameResolver, moduleData,
+                FirBuiltinAnnotationDeserializer(moduleData.session),
+                FirConstDeserializer(moduleData.session),
+                containerSource = null
+            ).memberDeserializer
+        }
+
+        private val lookup = moduleData.session.firCachesFactory.createCacheWithPostCompute(
+            { classId: ClassId, context: FirDeserializationContext? -> FirRegularClassSymbol(classId) to context }
+        ) { classId, symbol, parentContext ->
+            val classData = classDataFinder.findClassData(classId)!!
+            val classProto = classData.classProto
+
+            deserializeClassToSymbol(
+                classId, classProto, symbol, nameResolver, moduleData.session, moduleData,
+                null, kotlinScopeProvider, parentContext,
+                null,
+                origin = FirDeclarationOrigin.BuiltIns,
+                this::findAndDeserializeClass,
+            )
+        }
+
+        fun getClassLikeSymbolByFqName(classId: ClassId): FirRegularClassSymbol? =
+            findAndDeserializeClass(classId)
+
+        private fun findAndDeserializeClass(
+            classId: ClassId,
+            parentContext: FirDeserializationContext? = null,
+        ): FirRegularClassSymbol? {
+            val classIdExists = classId in classDataFinder.allClassIds
+            if (!classIdExists) return null
+            return lookup.getValue(classId, parentContext)
+        }
+
+        fun getTopLevelCallableSymbols(name: Name): List<FirCallableSymbol<*>> {
+            return getTopLevelFunctionSymbols(name)
+        }
+
+        fun getTopLevelFunctionSymbols(name: Name): List<FirNamedFunctionSymbol> {
+            return packageProto.`package`.functionList.filter { nameResolver.getName(it.name) == name }.map {
+                memberDeserializer.loadFunction(it).symbol
+            }
+        }
+    }
+}
+
+private data class BinaryVersionAndPackageFragment(
+    val version: BuiltInsBinaryVersion,
+    val packageFragment: ProtoBuf.PackageFragment,
+) {
+    companion object {
+        fun createFromStream(stream: InputStream): BinaryVersionAndPackageFragment {
+            val version = BuiltInsBinaryVersion.readFrom(stream)
+
+            if (!version.isCompatible()) {
+                // TODO: report a proper diagnostic
+                throw UnsupportedOperationException(
+                    "Kotlin built-in definition format version is not supported: " +
+                            "expected ${BuiltInsBinaryVersion.INSTANCE}, actual $version. " +
+                            "Please update Kotlin",
+                )
+            }
+
+            val packageFragment = ProtoBuf.PackageFragment.parseFrom(stream, BuiltInSerializerProtocol.extensionRegistry)
+            return BinaryVersionAndPackageFragment(version, packageFragment)
+        }
+    }
+}
+
+private class SyntheticFunctionalInterfaceCache(private val moduleData: FirModuleData, private val kotlinScopeProvider: FirKotlinScopeProvider) {
+    private val syntheticFunctionalInterfaceCache =
+        moduleData.session.firCachesFactory.createCache(::createSyntheticFunctionalInterface)
+
+    fun tryGetSyntheticFunctionalInterface(classId: ClassId): FirRegularClassSymbol? {
+        return syntheticFunctionalInterfaceCache.getValue(classId)
+    }
+
+    private fun createSyntheticFunctionalInterface(classId: ClassId): FirRegularClassSymbol? {
         return with(classId) {
             val className = relativeClassName.asString()
             val kind = FunctionClassKind.byClassNamePrefix(packageFqName, className) ?: return@with null
             val prefix = kind.classNamePrefix
             val arity = className.substring(prefix.length).toIntOrNull() ?: return null
-            syntheticFunctionalInterfaceSymbols.getOrPut(SyntheticFunctionalInterfaceSymbolKey(kind, arity)) {
+            FirRegularClassSymbol(classId).apply symbol@{
-                FirRegularClassSymbol(this).apply symbol@{
                 buildRegularClass klass@{
-                        moduleData = this@FirBuiltinSymbolProvider.moduleData
+                    moduleData = this@SyntheticFunctionalInterfaceCache.moduleData
                     origin = FirDeclarationOrigin.BuiltIns
                     name = relativeClassName.shortName()
                     status = FirResolvedDeclarationStatusImpl(
@@ -107,27 +219,27 @@ open class FirBuiltinSymbolProvider(
                     typeParameters.addAll(
                         (1..arity).map {
                             buildTypeParameter {
-                                    moduleData = this@FirBuiltinSymbolProvider.moduleData
+                                moduleData = this@SyntheticFunctionalInterfaceCache.moduleData
                                 resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
                                 origin = FirDeclarationOrigin.BuiltIns
                                 name = Name.identifier("P$it")
                                 symbol = FirTypeParameterSymbol()
                                 variance = Variance.IN_VARIANCE
                                 isReified = false
-                                    bounds += session.builtinTypes.nullableAnyType
+                                bounds += moduleData.session.builtinTypes.nullableAnyType
                             }
                         },
                     )
                     typeParameters.add(
                         buildTypeParameter {
-                                moduleData = this@FirBuiltinSymbolProvider.moduleData
+                            moduleData = this@SyntheticFunctionalInterfaceCache.moduleData
                             resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
                             origin = FirDeclarationOrigin.BuiltIns
                             name = Name.identifier("R")
                             symbol = FirTypeParameterSymbol()
                             variance = Variance.OUT_VARIANCE
                             isReified = false
-                                bounds += session.builtinTypes.nullableAnyType
+                            bounds += moduleData.session.builtinTypes.nullableAnyType
                         },
                     )
                     val name = OperatorNameConventions.INVOKE
@@ -196,7 +308,7 @@ open class FirBuiltinSymbolProvider(
                     }
                     addDeclaration(
                         buildSimpleFunction {
-                                moduleData = this@FirBuiltinSymbolProvider.moduleData
+                            moduleData = this@SyntheticFunctionalInterfaceCache.moduleData
                             resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
                             origin = FirDeclarationOrigin.BuiltIns
                             returnTypeRef = typeArguments.last()
@@ -209,7 +321,7 @@ open class FirBuiltinSymbolProvider(
                             valueParameters += typeArguments.dropLast(1).mapIndexed { index, typeArgument ->
                                 val parameterName = Name.identifier("p${index + 1}")
                                 buildValueParameter {
-                                        moduleData = this@FirBuiltinSymbolProvider.moduleData
+                                    moduleData = this@SyntheticFunctionalInterfaceCache.moduleData
                                     origin = FirDeclarationOrigin.BuiltIns
                                     resolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
                                     returnTypeRef = typeArgument
@@ -228,111 +340,7 @@ open class FirBuiltinSymbolProvider(
             }
         }
     }
-    }
+
 
     private fun FunctionClassKind.classId(arity: Int) = ClassId(packageFqName, numberedClassName(arity))
-
-    @FirSymbolProviderInternals
-    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
-        allPackageFragments[packageFqName]?.flatMapTo(destination) {
-            it.getTopLevelCallableSymbols(name)
-        }
-    }
-
-    @FirSymbolProviderInternals
-    override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
-        getTopLevelFunctionSymbolsToByPackageFragments(destination, packageFqName, name)
-    }
-
-    protected fun getTopLevelFunctionSymbolsToByPackageFragments(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
-        allPackageFragments[packageFqName]?.flatMapTo(destination) {
-            it.getTopLevelFunctionSymbols(name)
-        }
-    }
-
-    @FirSymbolProviderInternals
-    override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
-    }
-
-    private inner class BuiltInsPackageFragment(stream: InputStream, val fqName: FqName) {
-
-        private val binaryVersionAndPackageFragment = BinaryVersionAndPackageFragment.createFromStream(stream)
-
-        val version: BuiltInsBinaryVersion get() = binaryVersionAndPackageFragment.version
-        val packageProto: ProtoBuf.PackageFragment get() = binaryVersionAndPackageFragment.packageFragment
-
-        private val nameResolver = NameResolverImpl(packageProto.strings, packageProto.qualifiedNames)
-
-        val classDataFinder = ProtoBasedClassDataFinder(packageProto, nameResolver, version) { SourceElement.NO_SOURCE }
-
-        private val memberDeserializer by lazy {
-            FirDeserializationContext.createForPackage(
-                fqName, packageProto.`package`, nameResolver, moduleData,
-                FirBuiltinAnnotationDeserializer(session),
-                FirConstDeserializer(session),
-                containerSource = null
-            ).memberDeserializer
-        }
-
-        private val lookup = session.firCachesFactory.createCacheWithPostCompute(
-            { classId: ClassId, context: FirDeserializationContext? -> FirRegularClassSymbol(classId) to context }
-        ) { classId, symbol, parentContext ->
-            val classData = classDataFinder.findClassData(classId)!!
-            val classProto = classData.classProto
-
-            deserializeClassToSymbol(
-                classId, classProto, symbol, nameResolver, session,
-                moduleData,
-                    defaultAnnotationDeserializer =null, kotlinScopeProvider, parentContext,
-                containerSource =null,
-                origin = FirDeclarationOrigin.BuiltIns,
-                this::findAndDeserializeClass,
-            )
-        }
-
-        fun getClassLikeSymbolByFqName(classId: ClassId): FirRegularClassSymbol? =
-            findAndDeserializeClass(classId)
-
-        private fun findAndDeserializeClass(
-            classId: ClassId,
-            parentContext: FirDeserializationContext? = null,
-        ): FirRegularClassSymbol? {
-            val classIdExists = classId in classDataFinder.allClassIds
-            if (!classIdExists) return null
-            return lookup.getValue(classId, parentContext)
-        }
-
-        fun getTopLevelCallableSymbols(name: Name): List<FirCallableSymbol<*>> {
-            return getTopLevelFunctionSymbols(name)
-        }
-
-        fun getTopLevelFunctionSymbols(name: Name): List<FirNamedFunctionSymbol> {
-            return packageProto.`package`.functionList.filter { nameResolver.getName(it.name) == name }.map {
-                memberDeserializer.loadFunction(it).symbol
-            }
-        }
-    }
-}
-
-private data class BinaryVersionAndPackageFragment(
-    val version: BuiltInsBinaryVersion,
-    val packageFragment: ProtoBuf.PackageFragment,
-) {
-    companion object {
-        fun createFromStream(stream: InputStream): BinaryVersionAndPackageFragment {
-            val version = BuiltInsBinaryVersion.readFrom(stream)
-
-            if (!version.isCompatible()) {
-                // TODO: report a proper diagnostic
-                throw UnsupportedOperationException(
-                    "Kotlin built-in definition format version is not supported: " +
-                            "expected ${BuiltInsBinaryVersion.INSTANCE}, actual $version. " +
-                            "Please update Kotlin",
-                )
-            }
-
-            val packageFragment = ProtoBuf.PackageFragment.parseFrom(stream, BuiltInSerializerProtocol.extensionRegistry)
-           return BinaryVersionAndPackageFragment(version, packageFragment)
-        }
-    }
 }