Introduce ClassDescriptor.getSealedSubclasses

#KT-12795 In Progress

compiler/backend/src/org/jetbrains/kotlin/codegen/MutableClassDescriptor.java

@@ -33,7 +33,7 @@ import org.jetbrains.kotlin.types.TypeUtils;
 
 import java.util.*;
 
-public class MutableClassDescriptor extends ClassDescriptorBase implements ClassDescriptor {
+public class MutableClassDescriptor extends ClassDescriptorBase {
     private final ClassKind kind;
     private final boolean isInner;
 
@@ -71,6 +71,7 @@ public class MutableClassDescriptor extends ClassDescriptorBase implements Class
     }
 
     public void setModality(@NotNull Modality modality) {
+        assert modality != Modality.SEALED : "Implement getSealedSubclasses() for this class: " + getClass();
         this.modality = modality;
     }
 
@@ -180,6 +181,12 @@ public class MutableClassDescriptor extends ClassDescriptorBase implements Class
         return MemberScope.Empty.INSTANCE;
     }
 
+    @NotNull
+    @Override
+    public Collection<ClassDescriptor> getSealedSubclasses() {
+        return Collections.emptyList();
+    }
+
     @Override
     public String toString() {
         return DeclarationDescriptorImpl.toString(this);

compiler/frontend/src/org/jetbrains/kotlin/cfg/WhenChecker.kt

@@ -17,11 +17,9 @@
 package org.jetbrains.kotlin.cfg
 
 import com.intellij.psi.PsiElement
-import com.intellij.psi.tree.TokenSet
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ClassKind
-import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
@@ -36,11 +34,9 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils.isEnumEntry
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
 import org.jetbrains.kotlin.resolve.constants.CompileTimeConstant
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
-import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
-import org.jetbrains.kotlin.resolve.scopes.MemberScope
+import org.jetbrains.kotlin.resolve.descriptorUtil.computeSealedSubclasses
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
-import org.jetbrains.kotlin.types.isFlexible
 import java.util.*
 
 interface WhenMissingCase {
@@ -150,11 +146,12 @@ internal abstract class WhenOnClassExhaustivenessChecker : WhenExhaustivenessChe
 
     protected fun getMissingClassCases(
             whenExpression: KtWhenExpression,
-            memberDescriptors: Set<ClassDescriptor>,
+            subclasses: Set<ClassDescriptor>,
             context: BindingContext
     ): List<WhenMissingCase> {
         // when on empty enum / sealed is considered non-exhaustive, see test whenOnEmptySealed
-        if (memberDescriptors.isEmpty()) return listOf(UnknownMissingCase)
+        if (subclasses.isEmpty()) return listOf(UnknownMissingCase)
+
         val checkedDescriptors = LinkedHashSet<ClassDescriptor>()
         for (whenEntry in whenExpression.entries) {
             for (condition in whenEntry.conditions) {
@@ -182,7 +179,7 @@ internal abstract class WhenOnClassExhaustivenessChecker : WhenExhaustivenessChe
                 // Checks are important only for nested subclasses of the sealed class
                 // In additional, check without "is" is important only for objects
                 if (checkedDescriptor == null ||
-                    !memberDescriptors.contains(checkedDescriptor) ||
+                    !subclasses.contains(checkedDescriptor) ||
                     (condition is KtWhenConditionWithExpression &&
                      !DescriptorUtils.isObject(checkedDescriptor) &&
                      !DescriptorUtils.isEnumEntry(checkedDescriptor))) {
@@ -190,7 +187,7 @@ internal abstract class WhenOnClassExhaustivenessChecker : WhenExhaustivenessChe
                 }
                 if (negated) {
                     if (checkedDescriptors.contains(checkedDescriptor)) return listOf() // all members are already there
-                    checkedDescriptors.addAll(memberDescriptors)
+                    checkedDescriptors.addAll(subclasses)
                     checkedDescriptors.remove(checkedDescriptor)
                 }
                 else {
@@ -198,7 +195,7 @@ internal abstract class WhenOnClassExhaustivenessChecker : WhenExhaustivenessChe
                 }
             }
         }
-        return (memberDescriptors - checkedDescriptors).toList().map { ClassMissingCase(it) }
+        return (subclasses - checkedDescriptors).toList().map { ClassMissingCase(it) }
     }
 }
 
@@ -234,41 +231,15 @@ internal object WhenOnSealedExhaustivenessChecker : WhenOnClassExhaustivenessChe
         assert(DescriptorUtils.isSealedClass(subjectDescriptor)) {
             "isWhenOnSealedClassExhaustive should be called with a sealed class descriptor: $subjectDescriptor"
         }
-        val memberClassDescriptors = getNestedSubclasses(subjectDescriptor!!)
+        val subclasses = subjectDescriptor!!.sealedSubclasses
         // When on a sealed class without derived members is considered non-exhaustive (see test WhenOnEmptySealed)
-        return getMissingClassCases(expression, memberClassDescriptors, context) +
+        return getMissingClassCases(expression, subclasses.toSet(), context) +
                WhenOnNullableExhaustivenessChecker.getMissingCases(expression, context, nullable)
     }
 
     override fun isApplicable(subjectType: KotlinType): Boolean {
         return DescriptorUtils.isSealedClass(TypeUtils.getClassDescriptor(subjectType))
     }
-
-    internal fun getNestedSubclasses(baseDescriptor: ClassDescriptor): Set<ClassDescriptor> {
-        val memberClassDescriptors = LinkedHashSet<ClassDescriptor>()
-        collectNestedSubclasses(baseDescriptor, baseDescriptor, memberClassDescriptors)
-        return memberClassDescriptors
-    }
-
-    private fun collectNestedSubclasses(
-            baseDescriptor: ClassDescriptor,
-            currentDescriptor: ClassDescriptor,
-            subclasses: MutableSet<ClassDescriptor>
-    ) {
-        fun collectSubclasses(scope: MemberScope, collectNested: Boolean) {
-            for (descriptor in scope.getContributedDescriptors(DescriptorKindFilter.CLASSIFIERS)) {
-                if (descriptor is ClassDescriptor) {
-                    if (DescriptorUtils.isDirectSubclass(descriptor, baseDescriptor)) subclasses.add(descriptor)
-
-                    if (collectNested) collectNestedSubclasses(baseDescriptor, descriptor, subclasses)
-                }
-            }
-        }
-        if (currentDescriptor == baseDescriptor && DescriptorUtils.isTopLevelDeclaration(currentDescriptor)) {
-            collectSubclasses((currentDescriptor.containingDeclaration as PackageFragmentDescriptor).getMemberScope(), collectNested = false)
-        }
-        collectSubclasses(currentDescriptor.unsubstitutedInnerClassesScope, collectNested = true)
-    }
 }
 
 

compiler/frontend/src/org/jetbrains/kotlin/psi/synthetics/SyntheticClassOrObjectDescriptor.kt

@@ -83,6 +83,11 @@ class SyntheticClassOrObjectDescriptor(
     override fun getDeclaredTypeParameters() = emptyList<TypeParameterDescriptor>()
     override fun getStaticScope() = MemberScope.Empty
     override fun getUnsubstitutedMemberScope() = unsubstitutedMemberScope
+    override fun getSealedSubclasses() = emptyList<ClassDescriptor>()
+
+    init {
+        assert(modality != Modality.SEALED) { "Implement getSealedSubclasses() for this class: $javaClass" }
+    }
 
     override fun getDeclaredCallableMembers(): List<CallableMemberDescriptor> =
         DescriptorUtils.getAllDescriptors(unsubstitutedMemberScope).filterIsInstance<CallableMemberDescriptor>().filter {

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

@@ -39,6 +39,7 @@ import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.psi.psiUtil.KtPsiUtilKt;
 import org.jetbrains.kotlin.psi.synthetics.SyntheticClassOrObjectDescriptor;
 import org.jetbrains.kotlin.resolve.*;
+import org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilsKt;
 import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil;
 import org.jetbrains.kotlin.resolve.lazy.LazyClassContext;
 import org.jetbrains.kotlin.resolve.lazy.LazyEntity;
@@ -112,6 +113,8 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
 
     private final NotNullLazyValue<LexicalScope> scopeForInitializerResolution;
 
+    private final NotNullLazyValue<Collection<ClassDescriptor>> sealedSubclasses;
+
     public LazyClassDescriptor(
             @NotNull final LazyClassContext c,
             @NotNull DeclarationDescriptor containingDeclaration,
@@ -289,6 +292,14 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
                                                                                     classLikeInfo.getPrimaryConstructorParameters());
             }
         });
+
+        this.sealedSubclasses = storageManager.createLazyValue(new Function0<Collection<ClassDescriptor>>() {
+            @Override
+            public Collection<ClassDescriptor> invoke() {
+                // TODO: only consider classes from the same file, not the whole package fragment
+                return DescriptorUtilsKt.computeSealedSubclasses(LazyClassDescriptor.this);
+            }
+        });
     }
 
     @NotNull
@@ -549,6 +560,12 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
         return danglingAnnotations;
     }
 
+    @NotNull
+    @Override
+    public Collection<ClassDescriptor> getSealedSubclasses() {
+        return sealedSubclasses.invoke();
+    }
+
     @Override
     public String toString() {
         // not using descriptor render to preserve laziness

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/descriptors/KnownDescriptors.kt

@@ -41,6 +41,10 @@ open class KnownClassDescriptor(
         private val visibility: Visibility,
         override val annotations: Annotations
 ) : ClassDescriptor {
+    init {
+        assert(modality != Modality.SEALED) { "Implement getSealedSubclasses() for this class: $javaClass" }
+    }
+
     private lateinit var typeConstructor: TypeConstructor
     private lateinit var supertypes: List<KotlinType>
     private lateinit var defaultType: SimpleType
@@ -114,6 +118,7 @@ open class KnownClassDescriptor(
     override fun getContainingDeclaration(): DeclarationDescriptor = containingDeclaration
     override fun getDeclaredTypeParameters(): List<TypeParameterDescriptor> = declaredTypeParameters
     override fun getKind(): ClassKind = kind
+    override fun getSealedSubclasses(): Collection<ClassDescriptor> = emptyList()
 
     override fun getMemberScope(typeArguments: MutableList<out TypeProjection>): MemberScope = MemberScope.Empty
     override fun getMemberScope(typeSubstitution: TypeSubstitution): MemberScope = MemberScope.Empty
@@ -153,4 +158,3 @@ open class KnownClassDescriptor(
     override fun toString(): String =
             "KnownClassDescriptor($fqNameUnsafe)"
 }
-

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