Introduce ClassDescriptor.getSealedSubclasses

#KT-12795 In Progress

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaClassDescriptor.kt

@@ -135,6 +135,8 @@ class LazyJavaClassDescriptor(
 
     override fun getFunctionTypeForSamInterface(): SimpleType? = c.components.samConversionResolver.resolveFunctionTypeIfSamInterface(this)
 
+    override fun getSealedSubclasses(): Collection<ClassDescriptor> = emptyList()
+
     override fun toString() = "Lazy Java class ${this.fqNameUnsafe}"
 
     private inner class LazyJavaClassTypeConstructor : AbstractClassTypeConstructor(c.storageManager) {

core/descriptors/src/org/jetbrains/kotlin/builtins/functions/FunctionClassDescriptor.kt

@@ -106,7 +106,8 @@ class FunctionClassDescriptor(
     override fun isImpl() = false
     override fun isExternal() = false
     override val annotations: Annotations get() = Annotations.EMPTY
-    override fun getSource() = SourceElement.NO_SOURCE
+    override fun getSource(): SourceElement = SourceElement.NO_SOURCE
+    override fun getSealedSubclasses() = emptyList<ClassDescriptor>()
 
     override fun getDeclaredTypeParameters() = parameters
 

core/descriptors/src/org/jetbrains/kotlin/descriptors/ClassDescriptor.java

@@ -100,6 +100,13 @@ public interface ClassDescriptor extends ClassifierDescriptorWithTypeParameters,
     @NotNull
     List<TypeParameterDescriptor> getDeclaredTypeParameters();
 
+    /**
+     * @return direct subclasses of this class if it's a sealed class, empty list otherwise
+     */
+    @ReadOnly
+    @NotNull
+    Collection<ClassDescriptor> getSealedSubclasses();
+
     @NotNull
     @Override
     ClassDescriptor getOriginal();

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/ClassDescriptorImpl.java

@@ -51,6 +51,7 @@ public class ClassDescriptorImpl extends ClassDescriptorBase {
             boolean isExternal
     ) {
         super(LockBasedStorageManager.NO_LOCKS, containingDeclaration, name, source, isExternal);
+        assert modality != Modality.SEALED : "Implement getSealedSubclasses() for this class: " + getClass();
         this.modality = modality;
         this.kind = kind;
 
@@ -161,4 +162,10 @@ public class ClassDescriptorImpl extends ClassDescriptorBase {
     public List<TypeParameterDescriptor> getDeclaredTypeParameters() {
         return Collections.emptyList();
     }
+
+    @NotNull
+    @Override
+    public Collection<ClassDescriptor> getSealedSubclasses() {
+        return Collections.emptyList();
+    }
 }

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/EnumEntrySyntheticClassDescriptor.java

@@ -187,6 +187,12 @@ public class EnumEntrySyntheticClassDescriptor extends ClassDescriptorBase {
         return Collections.emptyList();
     }
 
+    @NotNull
+    @Override
+    public Collection<ClassDescriptor> getSealedSubclasses() {
+        return Collections.emptyList();
+    }
+
     private class EnumEntryScope extends MemberScopeImpl {
         private final MemoizedFunctionToNotNull<Name, Collection<SimpleFunctionDescriptor>> functions;
         private final MemoizedFunctionToNotNull<Name, Collection<PropertyDescriptor>> properties;

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/LazySubstitutingClassDescriptor.java

@@ -270,4 +270,10 @@ public class LazySubstitutingClassDescriptor implements ClassDescriptor {
         getSubstitutor();
         return declaredTypeParameters;
     }
+
+    @NotNull
+    @Override
+    public Collection<ClassDescriptor> getSealedSubclasses() {
+        return original.getSealedSubclasses();
+    }
 }

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

@@ -28,6 +28,8 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.FqNameUnsafe
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.constants.EnumValue
+import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
+import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeProjection
 import org.jetbrains.kotlin.types.TypeSubstitutor
@@ -351,3 +353,33 @@ 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): 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 = sealedClass.containingDeclaration
+    if (container is PackageFragmentDescriptor) {
+        collectSubclasses(container.getMemberScope(), collectNested = false)
+    }
+    collectSubclasses(sealedClass.unsubstitutedInnerClassesScope, collectNested = true)
+    return result
+}

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

@@ -110,6 +110,7 @@ class NotFoundClasses(private val storageManager: StorageManager, private val mo
         override fun getConstructors(): Collection<ClassConstructorDescriptor> = emptySet()
         override fun getUnsubstitutedPrimaryConstructor(): ClassConstructorDescriptor? = null
         override fun getCompanionObjectDescriptor(): ClassDescriptor? = null
+        override fun getSealedSubclasses(): Collection<ClassDescriptor> = emptyList()
 
         override fun toString() = "class $name (not found)"
     }

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

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.resolve.DescriptorFactory
 import org.jetbrains.kotlin.resolve.NonReportingOverrideStrategy
 import org.jetbrains.kotlin.resolve.OverridingUtil
 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
@@ -46,7 +47,7 @@ class DeserializedClassDescriptor(
         val classProto: ProtoBuf.Class,
         nameResolver: NameResolver,
         private val sourceElement: SourceElement
-) : ClassDescriptor, AbstractClassDescriptor(
+) : AbstractClassDescriptor(
         outerContext.storageManager,
         nameResolver.getClassId(classProto.fqName).shortClassName
 ) {
@@ -67,6 +68,7 @@ class DeserializedClassDescriptor(
     private val primaryConstructor = c.storageManager.createNullableLazyValue { computePrimaryConstructor() }
     private val constructors = c.storageManager.createLazyValue { computeConstructors() }
     private val companionObjectDescriptor = c.storageManager.createNullableLazyValue { computeCompanionObjectDescriptor() }
+    private val sealedSubclasses = c.storageManager.createLazyValue { computeSubclassesForSealedClass() }
 
     internal val thisAsProtoContainer: ProtoContainer.Class = ProtoContainer.Class(
             classProto, c.nameResolver, c.typeTable, sourceElement,
@@ -147,6 +149,13 @@ class DeserializedClassDescriptor(
     internal fun hasNestedClass(name: Name): Boolean =
             name in memberScope.classNames
 
+    private fun computeSubclassesForSealedClass(): Collection<ClassDescriptor> {
+        // TODO: store the list of subclasses to metadata
+        return computeSealedSubclasses(this)
+    }
+
+    override fun getSealedSubclasses() = sealedSubclasses()
+
     override fun toString() = "deserialized class $name" // not using descriptor render to preserve laziness
 
     override fun getSource() = sourceElement