KT-41346 Implement NoReorderImplementation without reordering

- See the documentation for `DeserializedMemberScopeHelper` for the full
info about the fix and the issue
- Add `preserveDeclarationsOrdering` setting to
`DeserializationConfiguration`
- Enable this setting in `DeserializerForClassfileDecompiler`
- Also, use `List` instead of `Collection` to pass members to
`DeserializedMemberScope`. It is done to emphasize that the order of
the members is important and may be used
- Review: https://jetbrains.team/p/kt/review/1627
- ^KT-41346 Fixed

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

@@ -30,5 +30,14 @@ interface DeserializationConfiguration {
     val releaseCoroutines: Boolean
         get() = false
 
+    /**
+     * We may want to preserve the order of the declarations the same as in the serialized object
+     * (for example, to later create a decompiled code with the original order of declarations).
+     *
+     * It is required to avoid PSI-Stub mismatch errors like in KT-41346.
+     */
+    val preserveDeclarationsOrdering: Boolean
+        get() = false
+
     object Default : DeserializationConfiguration
 }

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

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.resolve.MemberComparator
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.resolve.scopes.MemberScopeImpl
 import org.jetbrains.kotlin.serialization.deserialization.DeserializationContext
+import org.jetbrains.kotlin.serialization.deserialization.MemberDeserializer
 import org.jetbrains.kotlin.serialization.deserialization.getName
 import org.jetbrains.kotlin.storage.getValue
 import org.jetbrains.kotlin.utils.Printer
@@ -45,7 +46,7 @@ abstract class DeserializedMemberScope protected constructor(
     classNames: () -> Collection<Name>
 ) : MemberScopeImpl() {
 
-    private val impl: Implementation = OptimizedImplementation(functionList, propertyList, typeAliasList)
+    private val impl: Implementation = createImplementation(functionList, propertyList, typeAliasList)
 
     internal val classNames by c.storageManager.createLazyValue { classNames().toSet() }
 
@@ -163,6 +164,40 @@ abstract class DeserializedMemberScope protected constructor(
         p.println("}")
     }
 
+    /**
+     * This interface was introduced to fix KT-41346.
+     *
+     * The first implementation, [OptimizedImplementation], is more space-efficient and performant. It does not
+     * preserve the order of declarations in [addFunctionsAndPropertiesTo] though, and have to restore it manually. It is used
+     * in most situations when the [DeserializedMemberScope] is created.
+     *
+     * The second implementation, [NoReorderImplementation], is less efficient, but it keeps the descriptors
+     * in the same order as in serialized ProtoBuf objects in [addFunctionsAndPropertiesTo]. It should be used only when
+     * [org.jetbrains.kotlin.serialization.deserialization.DeserializationConfiguration.preserveDeclarationsOrdering] is
+     * set to `true`, which is done during decompilation from deserialized descriptors.
+     *
+     * The decompiled descriptors are used to build PSI, which is then compared with PSI built directly from classfiles and metadata.
+     *
+     * If the declarations in the first and the second PSI go in a different order, PSI-Stub mismatch error is raised.
+     *
+     * PSI from classfiles and metadata uses the same order of the declarations as in the serialized ProtoBuf objects.
+     * This order is dictated by [MemberComparator].
+     *
+     * [OptimizedImplementation] uses [MemberComparator.NameAndTypeMemberComparator] to restore the same order
+     * of the declarations as it should be in serialized objects. However, this does not always work (for example, when
+     * the Kotlin classes were obfuscated by ProGuard).
+     *
+     * ProGuard may rename some declarations in serialized objects, and then the comparator will reorder them based on their new names.
+     * This will lead to PSI-Stub mismatch error since the declarations are now differently ordered.
+     *
+     * To avoid this, we have [NoReorderImplementation] implementation. It performs no reordering of the declarations at
+     * all. Since it is less space-efficient, it is used only the scope is going to be used during decompilation.
+
+     * [createImplementation] is used to create the correct implementation of [Implementation].
+     *
+     * Both [OptimizedImplementation] and [NoReorderImplementation] are made inner classes to have
+     * access to protected `getNonDeclared*` functions.
+     */
     private interface Implementation {
         val functionNames: Set<Name>
         val variableNames: Set<Name>
@@ -180,6 +215,16 @@ abstract class DeserializedMemberScope protected constructor(
         )
     }
 
+    private fun createImplementation(
+        functionList: List<ProtoBuf.Function>,
+        propertyList: List<ProtoBuf.Property>,
+        typeAliasList: List<ProtoBuf.TypeAlias>
+    ): Implementation =
+        if (c.components.configuration.preserveDeclarationsOrdering)
+            NoReorderImplementation(functionList, propertyList, typeAliasList)
+        else
+            OptimizedImplementation(functionList, propertyList, typeAliasList)
+
     private inner class OptimizedImplementation(
         functionList: List<ProtoBuf.Function>,
         propertyList: List<ProtoBuf.Property>,
@@ -329,8 +374,138 @@ abstract class DeserializedMemberScope protected constructor(
                 }
             }
 
+            // We perform the sort just in case
             subResult.sortWith(MemberComparator.NameAndTypeMemberComparator.INSTANCE)
             result.addAll(subResult)
         }
     }
+
+    /**
+     * Take a note that [NoReorderImplementation] still adds non-declared members together with directly declared in class.
+     * This is not a problem for ordering, since during decompilation from descriptors those non-declared members are just ignored,
+     * and the declared members will be added to decompiled text in the proper (i.e. original) order.
+     */
+    private inner class NoReorderImplementation(
+        private val functionList: List<ProtoBuf.Function>,
+        private val propertyList: List<ProtoBuf.Property>,
+        typeAliasList: List<ProtoBuf.TypeAlias>
+    ) : Implementation {
+
+        private val typeAliasList = if (c.components.configuration.typeAliasesAllowed) typeAliasList else emptyList()
+
+        private val declaredFunctions: List<SimpleFunctionDescriptor>
+                by c.storageManager.createLazyValue { computeFunctions() }
+
+        private val declaredProperties: List<PropertyDescriptor>
+                by c.storageManager.createLazyValue { computeProperties() }
+
+        private val allTypeAliases: List<TypeAliasDescriptor>
+                by c.storageManager.createLazyValue { computeTypeAliases() }
+
+        private val allFunctions: List<SimpleFunctionDescriptor>
+                by c.storageManager.createLazyValue { declaredFunctions + computeAllNonDeclaredFunctions() }
+
+        private val allProperties: List<PropertyDescriptor>
+                by c.storageManager.createLazyValue { declaredProperties + computeAllNonDeclaredProperties() }
+
+        private val typeAliasesByName: Map<Name, TypeAliasDescriptor>
+                by c.storageManager.createLazyValue { allTypeAliases.associateBy { it.name } }
+
+        private val functionsByName: Map<Name, Collection<SimpleFunctionDescriptor>>
+                by c.storageManager.createLazyValue { allFunctions.groupBy { it.name } }
+
+        private val propertiesByName: Map<Name, Collection<PropertyDescriptor>>
+                by c.storageManager.createLazyValue { allProperties.groupBy { it.name } }
+
+        override val functionNames by c.storageManager.createLazyValue {
+            functionList.mapToNames { it.name } + getNonDeclaredFunctionNames()
+        }
+
+        override val variableNames by c.storageManager.createLazyValue {
+            propertyList.mapToNames { it.name } + getNonDeclaredVariableNames()
+        }
+
+        override val typeAliasNames: Set<Name>
+            get() = typeAliasList.mapToNames { it.name }
+
+        private fun computeFunctions(): List<SimpleFunctionDescriptor> =
+            functionList.mapWithDeserializer { loadFunction(it).takeIf(::isDeclaredFunctionAvailable) }
+
+        private fun computeProperties(): List<PropertyDescriptor> =
+            propertyList.mapWithDeserializer { loadProperty(it) }
+
+        private fun computeTypeAliases(): List<TypeAliasDescriptor> =
+            typeAliasList.mapWithDeserializer { loadTypeAlias(it) }
+
+        private fun computeAllNonDeclaredFunctions(): List<SimpleFunctionDescriptor> =
+            getNonDeclaredFunctionNames().flatMap { computeNonDeclaredFunctionsForName(it) }
+
+        private fun computeAllNonDeclaredProperties(): List<PropertyDescriptor> =
+            getNonDeclaredVariableNames().flatMap { computeNonDeclaredPropertiesForName(it) }
+
+        private fun computeNonDeclaredFunctionsForName(name: Name): List<SimpleFunctionDescriptor> =
+            computeNonDeclaredDescriptors(name, declaredFunctions, ::computeNonDeclaredFunctions)
+
+        private fun computeNonDeclaredPropertiesForName(name: Name): List<PropertyDescriptor> =
+            computeNonDeclaredDescriptors(name, declaredProperties, ::computeNonDeclaredProperties)
+
+        override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<SimpleFunctionDescriptor> {
+            if (name !in functionNames) return emptyList()
+            return functionsByName[name].orEmpty()
+        }
+
+        override fun getTypeAliasByName(name: Name): TypeAliasDescriptor? {
+            return typeAliasesByName[name]
+        }
+
+        override fun getContributedVariables(name: Name, location: LookupLocation): Collection<PropertyDescriptor> {
+            if (name !in variableNames) return emptyList()
+            return propertiesByName[name].orEmpty()
+        }
+
+        override fun addFunctionsAndPropertiesTo(
+            result: MutableCollection<DeclarationDescriptor>,
+            kindFilter: DescriptorKindFilter,
+            nameFilter: (Name) -> Boolean,
+            location: LookupLocation
+        ) {
+            if (kindFilter.acceptsKinds(DescriptorKindFilter.VARIABLES_MASK)) {
+                allProperties.filterTo(result) { nameFilter(it.name) }
+            }
+
+            if (kindFilter.acceptsKinds(DescriptorKindFilter.FUNCTIONS_MASK)) {
+                allFunctions.filterTo(result) { nameFilter(it.name) }
+            }
+        }
+
+        /**
+         * We have to collect non-declared properties in such non-pretty way because we don't want to change the contract of the
+         * [computeNonDeclaredProperties] and [computeNonDeclaredFunctions] methods, because we do not want any performance penalties.
+         *
+         * [computeNonDeclared] may only add elements to the end of [MutableList], otherwise this function would not work properly.
+         */
+        private inline fun <T : DeclarationDescriptor> computeNonDeclaredDescriptors(
+            name: Name,
+            declaredDescriptors: List<T>,
+            computeNonDeclared: (Name, MutableList<T>) -> Unit
+        ): List<T> {
+            val declaredDescriptorsWithSameName = declaredDescriptors.filterTo(mutableListOf()) { it.name == name }
+            val nonDeclaredPropertiesStartIndex = declaredDescriptorsWithSameName.size
+
+            computeNonDeclared(name, declaredDescriptorsWithSameName)
+
+            return declaredDescriptorsWithSameName.subList(nonDeclaredPropertiesStartIndex, declaredDescriptorsWithSameName.size)
+        }
+
+        private inline fun <T : MessageLite> List<T>.mapToNames(getNameIndex: (T) -> Int): Set<Name> {
+            // `mutableSetOf` returns `LinkedHashSet`, it is important to preserve the order of the declarations.
+            return mapTo(mutableSetOf()) { c.nameResolver.getName(getNameIndex(it)) }
+        }
+
+        private inline fun <T : MessageLite, K : MemberDescriptor> List<T>.mapWithDeserializer(
+            deserialize: MemberDeserializer.(T) -> K?
+        ): List<K> {
+            return mapNotNull { c.memberDeserializer.deserialize(it) }
+        }
+    }
 }

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

@@ -57,8 +57,8 @@ class DeserializerForClassfileDecompiler(
             BinaryClassAnnotationAndConstantLoaderImpl(moduleDescriptor, notFoundClasses, storageManager, classFinder)
 
         val configuration = object : DeserializationConfiguration {
-            override val readDeserializedContracts: Boolean
+            override val readDeserializedContracts: Boolean = true
-                get() = true
+            override val preserveDeclarationsOrdering: Boolean = true
         }
 
         deserializationComponents = DeserializationComponents(