[FE] Extract computation of sealed class inheritors into separate component

This is needed to provide more optimal provider in IDE plugin

compiler/frontend/src/org/jetbrains/kotlin/frontend/di/injection.kt

@@ -103,6 +103,7 @@ private fun StorageComponentContainer.configurePlatformIndependentComponents() {
     useImpl<ClassicTypeSystemContextForCS>()
     useImpl<ClassicConstraintSystemUtilContext>()
     useInstance(ProgressManagerBasedCancellationChecker)
+    useInstance(SealedClassInheritorsProviderImpl)
 }
 
 /**

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/LazyClassContext.kt

@@ -48,4 +48,5 @@ interface LazyClassContext {
     val kotlinTypeChecker: NewKotlinTypeChecker
     val samConversionResolver: SamConversionResolver
     val additionalClassPartsProvider: AdditionalClassPartsProvider
+    val sealedClassInheritorsProvider: SealedClassInheritorsProvider
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/ResolveSession.java

@@ -85,6 +85,7 @@ public class ResolveSession implements KotlinCodeAnalyzer, LazyClassContext {
     private WrappedTypeFactory wrappedTypeFactory;
     private PlatformDiagnosticSuppressor platformDiagnosticSuppressor;
     private SamConversionResolver samConversionResolver;
+    private SealedClassInheritorsProvider sealedClassInheritorsProvider;
 
     private AdditionalClassPartsProvider additionalClassPartsProvider;
 
@@ -164,6 +165,11 @@ public class ResolveSession implements KotlinCodeAnalyzer, LazyClassContext {
         this.additionalClassPartsProvider = additionalClassPartsProvider;
     }
 
+    @Inject
+    public void setSealedClassInheritorsProvider(@NotNull SealedClassInheritorsProvider sealedClassInheritorsProvider) {
+        this.sealedClassInheritorsProvider = sealedClassInheritorsProvider;
+    }
+
     // Only calls from injectors expected
     @Deprecated
     public ResolveSession(
@@ -513,4 +519,10 @@ public class ResolveSession implements KotlinCodeAnalyzer, LazyClassContext {
     public AdditionalClassPartsProvider getAdditionalClassPartsProvider() {
         return additionalClassPartsProvider;
     }
+
+    @NotNull
+    @Override
+    public SealedClassInheritorsProvider getSealedClassInheritorsProvider() {
+        return sealedClassInheritorsProvider;
+    }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/LazyClassDescriptor.java

@@ -270,9 +270,8 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
                 )
         );
 
-        // TODO: only consider classes from the same file, not the whole package fragment
         boolean freedomForSealedInterfacesSupported = c.getLanguageVersionSettings().supportsFeature(LanguageFeature.FreedomForSealedClasses);
-        this.sealedSubclasses = storageManager.createLazyValue(() -> DescriptorUtilsKt.computeSealedSubclasses(this, freedomForSealedInterfacesSupported));
+        this.sealedSubclasses = storageManager.createLazyValue(() -> c.getSealedClassInheritorsProvider().computeSealedSubclasses(this, freedomForSealedInterfacesSupported));
     }
 
     private static boolean isIllegalInner(@NotNull DeclarationDescriptor descriptor) {

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/LocalClassifierAnalyzer.kt

@@ -71,6 +71,7 @@ class LocalClassifierAnalyzer(
     private val kotlinTypeChecker: NewKotlinTypeChecker,
     private val samConversionResolver: SamConversionResolver,
     private val additionalClassPartsProvider: AdditionalClassPartsProvider,
+    private val sealedClassInheritorsProvider: SealedClassInheritorsProvider
 ) {
     fun processClassOrObject(
         scope: LexicalWritableScope?,
@@ -107,6 +108,7 @@ class LocalClassifierAnalyzer(
                 kotlinTypeChecker,
                 samConversionResolver,
                 additionalClassPartsProvider,
+                sealedClassInheritorsProvider
             ),
             analyzerServices
         )
@@ -138,6 +140,7 @@ class LocalClassDescriptorHolder(
     val kotlinTypeChecker: NewKotlinTypeChecker,
     val samConversionResolver: SamConversionResolver,
     val additionalClassPartsProvider: AdditionalClassPartsProvider,
+    val sealedClassInheritorsProvider: SealedClassInheritorsProvider
 ) {
     // We do not need to synchronize here, because this code is used strictly from one thread
     private var classDescriptor: ClassDescriptor? = null
@@ -181,6 +184,7 @@ class LocalClassDescriptorHolder(
                     override val samConversionResolver: SamConversionResolver = this@LocalClassDescriptorHolder.samConversionResolver
                     override val additionalClassPartsProvider: AdditionalClassPartsProvider =
                         this@LocalClassDescriptorHolder.additionalClassPartsProvider
+                    override val sealedClassInheritorsProvider: SealedClassInheritorsProvider = this@LocalClassDescriptorHolder.sealedClassInheritorsProvider
                 },
                 containingDeclaration,
                 classOrObject.nameAsSafeName,

core/descriptors/src/org/jetbrains/kotlin/resolve/DescriptorUtils.kt

@@ -374,43 +374,6 @@ fun ClassifierDescriptor.getAllSuperClassifiers(): Sequence<ClassifierDescriptor
     return doGetAllSuperClassesAndInterfaces()
 }
 
-// Note this is a generic and slow implementation which would work almost for any subclass of ClassDescriptor.
-// Please avoid using it in new code.
-// TODO: do something more clever instead at call sites of this function
-fun computeSealedSubclasses(sealedClass: ClassDescriptor, freedomForSealedInterfacesSupported: Boolean): Collection<ClassDescriptor> {
-    if (sealedClass.modality != Modality.SEALED) return emptyList()
-
-    val result = linkedSetOf<ClassDescriptor>()
-
-    fun collectSubclasses(scope: MemberScope, collectNested: Boolean) {
-        for (descriptor in scope.getContributedDescriptors(DescriptorKindFilter.CLASSIFIERS)) {
-            if (descriptor !is ClassDescriptor) continue
-
-            if (DescriptorUtils.isDirectSubclass(descriptor, sealedClass)) {
-                result.add(descriptor)
-            }
-
-            if (collectNested) {
-                collectSubclasses(descriptor.unsubstitutedInnerClassesScope, collectNested)
-            }
-        }
-    }
-
-    val container = if (!freedomForSealedInterfacesSupported) {
-        sealedClass.containingDeclaration
-    } else {
-        sealedClass.parents.firstOrNull { it is PackageFragmentDescriptor }
-    }
-    if (container is PackageFragmentDescriptor) {
-        collectSubclasses(
-            container.getMemberScope(),
-            collectNested = freedomForSealedInterfacesSupported
-        )
-    }
-    collectSubclasses(sealedClass.unsubstitutedInnerClassesScope, collectNested = true)
-    return result
-}
-
 fun DeclarationDescriptor.isPublishedApi(): Boolean {
     val descriptor = if (this is CallableMemberDescriptor) DescriptorUtils.getDirectMember(this) else this
     return descriptor.annotations.hasAnnotation(StandardNames.FqNames.publishedApi)

core/descriptors/src/org/jetbrains/kotlin/resolve/SealedClassInheritorsProvider.kt

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.resolve
+
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
+import org.jetbrains.kotlin.resolve.descriptorUtil.parents
+import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
+import org.jetbrains.kotlin.resolve.scopes.MemberScope
+
+abstract class SealedClassInheritorsProvider {
+    abstract fun computeSealedSubclasses(
+        sealedClass: ClassDescriptor,
+        freedomForSealedInterfacesSupported: Boolean
+    ): Collection<ClassDescriptor>
+}
+
+object SealedClassInheritorsProviderImpl : SealedClassInheritorsProvider() {
+    // Note this is a generic and slow implementation which would work almost for any subclass of ClassDescriptor.
+    // Please avoid using it in new code.
+    // TODO: do something more clever instead at call sites of this function
+    override fun computeSealedSubclasses(
+        sealedClass: ClassDescriptor,
+        freedomForSealedInterfacesSupported: Boolean
+    ): Collection<ClassDescriptor> {
+        if (sealedClass.modality != Modality.SEALED) return emptyList()
+
+        val result = linkedSetOf<ClassDescriptor>()
+
+        fun collectSubclasses(scope: MemberScope, collectNested: Boolean) {
+            for (descriptor in scope.getContributedDescriptors(DescriptorKindFilter.CLASSIFIERS)) {
+                if (descriptor !is ClassDescriptor) continue
+
+                if (DescriptorUtils.isDirectSubclass(descriptor, sealedClass)) {
+                    result.add(descriptor)
+                }
+
+                if (collectNested) {
+                    collectSubclasses(descriptor.unsubstitutedInnerClassesScope, collectNested)
+                }
+            }
+        }
+
+        val container = if (!freedomForSealedInterfacesSupported) {
+            sealedClass.containingDeclaration
+        } else {
+            sealedClass.parents.firstOrNull { it is PackageFragmentDescriptor }
+        }
+        if (container is PackageFragmentDescriptor) {
+            collectSubclasses(
+                container.getMemberScope(),
+                collectNested = freedomForSealedInterfacesSupported
+            )
+        }
+        collectSubclasses(sealedClass.unsubstitutedInnerClassesScope, collectNested = true)
+        return result
+    }
+
+}

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

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.deserialization.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.protobuf.ExtensionRegistryLite
+import org.jetbrains.kotlin.resolve.SealedClassInheritorsProvider
 import org.jetbrains.kotlin.resolve.constants.ConstantValue
 import org.jetbrains.kotlin.resolve.sam.SamConversionResolver
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource

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

@@ -18,8 +18,8 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorFactory
 import org.jetbrains.kotlin.resolve.NonReportingOverrideStrategy
 import org.jetbrains.kotlin.resolve.OverridingUtil
+import org.jetbrains.kotlin.resolve.SealedClassInheritorsProviderImpl
 import org.jetbrains.kotlin.resolve.descriptorUtil.classId
-import org.jetbrains.kotlin.resolve.descriptorUtil.computeSealedSubclasses
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.resolve.scopes.StaticScopeForKotlinEnum
@@ -165,7 +165,7 @@ class DeserializedClassDescriptor(
         }
 
         // This is needed because classes compiled with Kotlin 1.0 did not contain the sealed_subclass_fq_name field
-        return computeSealedSubclasses(this, freedomForSealedInterfacesSupported = false)
+        return SealedClassInheritorsProviderImpl.computeSealedSubclasses(this, freedomForSealedInterfacesSupported = false)
     }
 
     override fun getSealedSubclasses() = sealedSubclasses()

native/commonizer/src/org/jetbrains/kotlin/descriptors/commonizer/builder/CommonizedClassDescriptor.kt

@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.descriptors.impl.ClassDescriptorBase
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorFactory.createPrimaryConstructorForObject
-import org.jetbrains.kotlin.resolve.descriptorUtil.computeSealedSubclasses
+import org.jetbrains.kotlin.resolve.SealedClassInheritorsProviderImpl
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.resolve.scopes.StaticScopeForKotlinEnum
@@ -54,7 +54,7 @@ class CommonizedClassDescriptor(
     private val typeConstructor = CommonizedClassTypeConstructor(targetComponents, cirSupertypes)
     private val sealedSubclasses = targetComponents.storageManager.createLazyValue {
         // TODO: pass proper language version settings
-        computeSealedSubclasses(this, freedomForSealedInterfacesSupported = false)
+        SealedClassInheritorsProviderImpl.computeSealedSubclasses(this, freedomForSealedInterfacesSupported = false)
     }
 
     private val declaredTypeParametersAndTypeParameterResolver = targetComponents.storageManager.createLazyValue {