Avoid retaining all protobuf objects in DeserializedMemberScope

Object model for deserialized protobuf isn't very cheap, so
loading just our stdlib's packages several times (as we do in the IDE)
might lead to ~ 96 MB of retained memory just by these objects.

While in fact many of them remain unused

 #KT-21517 Fixed

core/descriptors.jvm/src/org/jetbrains/kotlin/builtins/JvmBuiltInsPackageFragmentProvider.kt

@@ -53,7 +53,8 @@ class JvmBuiltInsPackageFragmentProvider(
                 ),
                 notFoundClasses,
                 ContractDeserializer.DEFAULT,
-                additionalClassPartsProvider, platformDependentDeclarationFilter
+                additionalClassPartsProvider, platformDependentDeclarationFilter,
+                BuiltInSerializerProtocol.extensionRegistry
         )
     }
 

core/descriptors.jvm/src/org/jetbrains/kotlin/load/kotlin/DeserializationComponentsForJava.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.load.java.lazy.LazyJavaPackageFragmentProvider
 import org.jetbrains.kotlin.platform.JvmBuiltIns
 import org.jetbrains.kotlin.serialization.deserialization.*
+import org.jetbrains.kotlin.serialization.jvm.JvmProtoBufUtil
 import org.jetbrains.kotlin.storage.StorageManager
 
 // This class is needed only for easier injection: exact types of needed components are specified in the constructor here.
@@ -50,7 +51,8 @@ class DeserializationComponentsForJava(
                 LocalClassifierTypeSettings.Default, errorReporter, lookupTracker, JavaFlexibleTypeDeserializer,
                 emptyList(), notFoundClasses, contractDeserializer,
                 additionalClassPartsProvider = jvmBuiltIns?.settings ?: AdditionalClassPartsProvider.None,
-                platformDependentDeclarationFilter = jvmBuiltIns?.settings ?: PlatformDependentDeclarationFilter.NoPlatformDependent
+                platformDependentDeclarationFilter = jvmBuiltIns?.settings ?: PlatformDependentDeclarationFilter.NoPlatformDependent,
+                extensionRegistryLite = JvmProtoBufUtil.EXTENSION_REGISTRY
         )
     }
 }

core/deserialization/src/org/jetbrains/kotlin/builtins/BuiltInsLoaderImpl.kt

@@ -84,7 +84,8 @@ class BuiltInsLoaderImpl : BuiltInsLoader {
                 notFoundClasses,
                 ContractDeserializer.DEFAULT,
                 additionalClassPartsProvider,
-                platformDependentDeclarationFilter
+                platformDependentDeclarationFilter,
+                BuiltInSerializerProtocol.extensionRegistry
         )
 
         for (packageFragment in packageFragments) {

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/MetadataPackageFragmentProvider.kt

@@ -24,7 +24,6 @@ import org.jetbrains.kotlin.descriptors.PackagePartProvider
 import org.jetbrains.kotlin.descriptors.SourceElement
 import org.jetbrains.kotlin.descriptors.deserialization.AdditionalClassPartsProvider
 import org.jetbrains.kotlin.descriptors.deserialization.PlatformDependentDeclarationFilter
-import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -60,7 +59,8 @@ class MetadataPackageFragmentProvider(
                 emptyList(),
                 notFoundClasses,
                 ContractDeserializer.DEFAULT,
-                AdditionalClassPartsProvider.None, PlatformDependentDeclarationFilter.All
+                AdditionalClassPartsProvider.None, PlatformDependentDeclarationFilter.All,
+                BuiltInSerializerProtocol.extensionRegistry
         )
     }
 

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/context.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.descriptors.deserialization.ClassDescriptorFactory
 import org.jetbrains.kotlin.descriptors.deserialization.PlatformDependentDeclarationFilter
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite
 import org.jetbrains.kotlin.resolve.constants.ConstantValue
 import org.jetbrains.kotlin.serialization.ProtoBuf
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
@@ -45,7 +46,8 @@ class DeserializationComponents(
         val notFoundClasses: NotFoundClasses,
         val contractDeserializer: ContractDeserializer,
         val additionalClassPartsProvider: AdditionalClassPartsProvider = AdditionalClassPartsProvider.None,
-        val platformDependentDeclarationFilter: PlatformDependentDeclarationFilter = PlatformDependentDeclarationFilter.All
+        val platformDependentDeclarationFilter: PlatformDependentDeclarationFilter = PlatformDependentDeclarationFilter.All,
+        val extensionRegistryLite: ExtensionRegistryLite
 ) {
     val classDeserializer: ClassDeserializer = ClassDeserializer(this)
 

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/descriptors/DeserializedMemberScope.kt

@@ -20,7 +20,9 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.protobuf.AbstractMessageLite
 import org.jetbrains.kotlin.protobuf.MessageLite
+import org.jetbrains.kotlin.protobuf.Parser
 import org.jetbrains.kotlin.resolve.MemberComparator
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.resolve.scopes.MemberScopeImpl
@@ -30,6 +32,8 @@ import org.jetbrains.kotlin.storage.getValue
 import org.jetbrains.kotlin.utils.Printer
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.compact
+import java.io.ByteArrayInputStream
+import java.io.ByteArrayOutputStream
 import java.util.*
 
 abstract class DeserializedMemberScope protected constructor(
@@ -40,20 +44,22 @@ abstract class DeserializedMemberScope protected constructor(
         classNames: () -> Collection<Name>
 ) : MemberScopeImpl() {
 
-    private val functionProtos by
+    private val functionProtosBytes = functionList.groupByName { it.name }.packToByteArray()
-            c.storageManager.createLazyValue {
+
-                functionList.groupByName { it.name }
+    private val propertyProtosBytes = propertyList.groupByName { it.name }.packToByteArray()
-            }
+
-    private val propertyProtos by
+    private val typeAliasBytes =
-            c.storageManager.createLazyValue {
+        if (c.components.configuration.typeAliasesAllowed)
-                propertyList.groupByName { it.name }
+            typeAliasList.groupByName { it.name }.packToByteArray()
-            }
+        else
-    private val typeAliasProtos by
+            emptyMap()
-            c.storageManager.createLazyValue {
+
-                if (c.components.configuration.typeAliasesAllowed)
+    private fun Map<Name, Collection<AbstractMessageLite>>.packToByteArray(): Map<Name, ByteArray> =
-                    typeAliasList.groupByName { it.name }
+        mapValues { entry ->
-                else emptyMap()
+            val byteArrayOutputStream = ByteArrayOutputStream()
-            }
+            entry.value.map { proto -> proto.writeDelimitedTo(byteArrayOutputStream) }
+            byteArrayOutputStream.toByteArray()
+        }
 
     private val functions =
             c.storageManager.createMemoizedFunction<Name, Collection<SimpleFunctionDescriptor>> { computeFunctions(it) }
@@ -63,14 +69,14 @@ abstract class DeserializedMemberScope protected constructor(
             c.storageManager.createMemoizedFunctionWithNullableValues<Name, TypeAliasDescriptor> { createTypeAlias(it) }
 
     private val functionNamesLazy by c.storageManager.createLazyValue {
-        functionProtos.keys + getNonDeclaredFunctionNames()
+        functionProtosBytes.keys + getNonDeclaredFunctionNames()
     }
 
     private val variableNamesLazy by c.storageManager.createLazyValue {
-        propertyProtos.keys + getNonDeclaredVariableNames()
+        propertyProtosBytes.keys + getNonDeclaredVariableNames()
     }
 
-    private val typeAliasNames: Set<Name> get() = typeAliasProtos.keys
+    private val typeAliasNames: Set<Name> get() = typeAliasBytes.keys
 
     internal val classNames by c.storageManager.createLazyValue { classNames().toSet() }
 
@@ -87,21 +93,37 @@ abstract class DeserializedMemberScope protected constructor(
     ) = groupBy { c.nameResolver.getName(getNameIndex(it)) }
 
     private fun computeFunctions(name: Name) =
-            computeDescriptors(
+        computeDescriptors(
-                    name,
+            name,
-                    functionProtos,
+            functionProtosBytes,
-                    { c.memberDeserializer.loadFunction(it) },
+            ProtoBuf.Function.PARSER,
-                    { computeNonDeclaredFunctions(name, it)}
+            { c.memberDeserializer.loadFunction(it) },
-            )
+            { computeNonDeclaredFunctions(name, it) }
+        )
 
-    inline private fun <M : MessageLite, D : DeclarationDescriptor> computeDescriptors(
+    private inline fun <M : MessageLite, D : DeclarationDescriptor> computeDescriptors(
-            name: Name,
+        name: Name,
-            protosByName: Map<Name, Collection<M>>,
+        bytesByName: Map<Name, ByteArray>,
-            factory: (M) -> D,
+        parser: Parser<M>,
-            computeNonDeclared: (MutableCollection<D>) -> Unit
+        factory: (M) -> D,
+        computeNonDeclared: (MutableCollection<D>) -> Unit
+    ): Collection<D> =
+        computeDescriptors(
+            bytesByName[name]?.let {
+                val inputStream = ByteArrayInputStream(it)
+                generateSequence {
+                    parser.parseDelimitedFrom(inputStream, c.components.extensionRegistryLite)
+                }.toList()
+            } ?: emptyList(),
+            factory,
+            computeNonDeclared
+        )
+
+    private inline fun <M : MessageLite, D : DeclarationDescriptor> computeDescriptors(
+        protos: Collection<M>,
+        factory: (M) -> D,
+        computeNonDeclared: (MutableCollection<D>) -> Unit
     ): Collection<D> {
-        val protos = protosByName[name].orEmpty()
-
         val descriptors = protos.mapTo(arrayListOf(), factory)
 
         computeNonDeclared(descriptors)
@@ -117,20 +139,22 @@ abstract class DeserializedMemberScope protected constructor(
     }
 
     private fun computeProperties(name: Name) =
-            computeDescriptors(
+        computeDescriptors(
-                    name,
+            name,
-                    propertyProtos,
+            propertyProtosBytes,
-                    { c.memberDeserializer.loadProperty(it) },
+            ProtoBuf.Property.PARSER,
-                    { computeNonDeclaredProperties(name, it) }
+            { c.memberDeserializer.loadProperty(it) },
-            )
+            { computeNonDeclaredProperties(name, it) }
+        )
 
     protected open fun computeNonDeclaredProperties(name: Name, descriptors: MutableCollection<PropertyDescriptor>) {
     }
 
-    private fun createTypeAlias(name: Name) =
+    private fun createTypeAlias(name: Name): TypeAliasDescriptor? {
-            typeAliasProtos[name]?.singleOrNull()?.let {
+        val byteArray = typeAliasBytes[name] ?: return null
-                c.memberDeserializer.loadTypeAlias(it)
+        val proto = ProtoBuf.TypeAlias.parseDelimitedFrom(ByteArrayInputStream(byteArray)) ?: return null
-            }
+        return c.memberDeserializer.loadTypeAlias(proto)
+    }
 
     override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> {
         if (name !in getVariableNames()) return emptyList()

idea/idea-analysis/src/org/jetbrains/kotlin/idea/decompiler/classFile/DeserializerForClassfileDecompiler.kt

@@ -71,7 +71,8 @@ class DeserializerForClassfileDecompiler(
                 storageManager, moduleDescriptor, DeserializationConfiguration.Default, classDataFinder, annotationAndConstantLoader,
                 packageFragmentProvider, ResolveEverythingToKotlinAnyLocalClassifierResolver(builtIns), LoggingErrorReporter(LOG),
                 LookupTracker.DO_NOTHING, JavaFlexibleTypeDeserializer, emptyList(), notFoundClasses,
-                ContractDeserializer.DEFAULT
+                ContractDeserializer.DEFAULT,
+                extensionRegistryLite = JvmProtoBufUtil.EXTENSION_REGISTRY
         )
     }
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/decompiler/common/KotlinMetadataDeserializerForDecompiler.kt

@@ -52,7 +52,8 @@ class KotlinMetadataDeserializerForDecompiler(
                 AnnotationAndConstantLoaderImpl(moduleDescriptor, notFoundClasses, serializerProtocol), packageFragmentProvider,
                 ResolveEverythingToKotlinAnyLocalClassifierResolver(builtIns), LoggingErrorReporter(LOG),
                 LookupTracker.DO_NOTHING, flexibleTypeDeserializer, emptyList(), notFoundClasses,
-                ContractDeserializer.DEFAULT
+                ContractDeserializer.DEFAULT,
+                extensionRegistryLite = serializerProtocol.extensionRegistry
         )
     }
 

js/js.serializer/src/org/jetbrains/kotlin/serialization/js/kotlinJavascriptPackageFragmentProvider.kt

@@ -73,7 +73,8 @@ fun createKotlinJavascriptPackageFragmentProvider(
             emptyList(),
             notFoundClasses,
             ContractDeserializerImpl(configuration),
-            platformDependentDeclarationFilter = PlatformDependentDeclarationFilter.NoPlatformDependent
+            platformDependentDeclarationFilter = PlatformDependentDeclarationFilter.NoPlatformDependent,
+            extensionRegistryLite = JsSerializerProtocol.extensionRegistry
     )
 
     for (packageFragment in packageFragments.filterIsInstance<KotlinJavascriptPackageFragment>()) {