[AA Standalone] Consider type aliases in direct inheritors search

A type alias may still be inherited from. For example:

```
sealed class MyClass

typealias T = MyClass

class Inheritor : T() // `Inheritor` is a direct inheritor of `MyClass`.
```

The index is a simplified version of the IDE's
`KotlinTypeAliasByExpansionShortNameIndex`, but it should be sufficient
for virtually all cases.

^KT-66013

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticDeclarationIndex.kt

@@ -22,4 +22,28 @@ internal class KotlinStaticDeclarationIndex {
      * Allows quickly finding [KtClassOrObject]s which have a given simple name as a supertype. The map may contain local classes as well.
      */
     internal val classesBySupertypeName: MutableMap<Name, MutableSet<KtClassOrObject>> = mutableMapOf()
+
+    /**
+     * Maps a simple name `N` to type aliases `A` in whose definition `N` occurs as the topmost user type, which is a prerequisite for other
+     * classes inheriting from `N` by referring to `A`. Does not support function types (e.g. `Function1`).
+     *
+     * There is no guarantee that the type alias can be inherited from. For example, if its expanded type is final, the type alias is not
+     * inheritable. The resulting type alias `A` may also occur in the expanded type of another type alias (which may also be inheritable),
+     * so the index may need to be followed transitively.
+     *
+     * The index is used to find direct class inheritors.
+     *
+     * ### Example
+     *
+     * ```
+     * abstract class C
+     *
+     * typealias A = C
+     *
+     * class X : A()
+     * ```
+     *
+     * The index contains the following entry: `"C" -> A`.
+     */
+    internal val inheritableTypeAliasesByAliasedName: MutableMap<Name, MutableSet<KtTypeAlias>> = mutableMapOf()
 }

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticDeclarationProvider.kt

@@ -35,11 +35,13 @@ import org.jetbrains.kotlin.fileClasses.javaFileFacadeFqName
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getImportedSimpleNameByImportAlias
 import org.jetbrains.kotlin.psi.psiUtil.getSuperNames
 import org.jetbrains.kotlin.psi.stubs.KotlinClassOrObjectStub
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.psi.stubs.impl.*
 import org.jetbrains.kotlin.serialization.deserialization.builtins.BuiltInSerializerProtocol
+import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.flattenTo
 import java.util.concurrent.ConcurrentHashMap
 
@@ -209,7 +211,7 @@ public class KotlinStaticDeclarationProviderFactory(
         }
 
         override fun visitTypeAlias(typeAlias: KtTypeAlias) {
-            addToTypeAliasMap(typeAlias)
+            indexTypeAlias(typeAlias)
             super.visitTypeAlias(typeAlias)
         }
 
@@ -258,6 +260,11 @@ public class KotlinStaticDeclarationProviderFactory(
         }
     }
 
+    private fun indexTypeAlias(typeAlias: KtTypeAlias) {
+        addToTypeAliasMap(typeAlias)
+        indexTypeAliasDefinition(typeAlias)
+    }
+
     private fun addToTypeAliasMap(typeAlias: KtTypeAlias) {
         typeAlias.getClassId()?.let { classId ->
             index.typeAliasMap.computeIfAbsent(classId.packageFqName) {
@@ -266,6 +273,44 @@ public class KotlinStaticDeclarationProviderFactory(
         }
     }
 
+    private fun indexTypeAliasDefinition(typeAlias: KtTypeAlias) {
+        val typeElement = typeAlias.getTypeReference()?.typeElement ?: return
+
+        findInheritableSimpleNames(typeElement).forEach { expandedName ->
+            index.inheritableTypeAliasesByAliasedName
+                .computeIfAbsent(Name.identifier(expandedName)) { mutableSetOf() }
+                .add(typeAlias)
+        }
+    }
+
+    /**
+     * This is a simplified version of `KtTypeElement.index()` from the IDE. If we need to move more indexing code to Standalone, we should
+     * consider moving more code from the IDE to the Analysis API.
+     *
+     * @see KotlinStaticDeclarationIndex.inheritableTypeAliasesByAliasedName
+     */
+    private fun findInheritableSimpleNames(typeElement: KtTypeElement): List<String> {
+        return when (typeElement) {
+            is KtUserType -> {
+                val referenceName = typeElement.referencedName ?: return emptyList()
+
+                buildList {
+                    add(referenceName)
+
+                    val ktFile = typeElement.containingKtFile
+                    if (!ktFile.isCompiled) {
+                        addIfNotNull(getImportedSimpleNameByImportAlias(typeElement.containingKtFile, referenceName))
+                    }
+                }
+            }
+
+            // `typealias T = A?` is inheritable.
+            is KtNullableType -> typeElement.innerType?.let(::findInheritableSimpleNames) ?: emptyList()
+
+            else -> emptyList()
+        }
+    }
+
     private fun addToFunctionMap(function: KtNamedFunction) {
         if (!function.isTopLevel) return
         val packageFqName = (function.parent as KtFile).packageFqName
@@ -330,7 +375,7 @@ public class KotlinStaticDeclarationProviderFactory(
                 // member functions and properties
                 stub.childrenStubs.forEach(::indexStub)
             }
-            is KotlinTypeAliasStubImpl -> addToTypeAliasMap(stub.psi)
+            is KotlinTypeAliasStubImpl -> indexTypeAlias(stub.psi)
             is KotlinFunctionStubImpl -> addToFunctionMap(stub.psi)
             is KotlinPropertyStubImpl -> addToPropertyMap(stub.psi)
             is KotlinPlaceHolderStubImpl -> {
@@ -416,6 +461,9 @@ public class KotlinStaticDeclarationProviderFactory(
 
     public fun getDirectInheritorCandidates(baseClassName: Name): Set<KtClassOrObject> =
         index.classesBySupertypeName[baseClassName].orEmpty()
+
+    public fun getInheritableTypeAliases(aliasedName: Name): Set<KtTypeAlias> =
+        index.inheritableTypeAliasesByAliasedName[aliasedName].orEmpty()
 }
 
 /**

analysis/analysis-api-standalone/analysis-api-fir-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/providers/KotlinStandaloneDirectInheritorsProvider.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.psiUtil.contains
@@ -42,7 +43,11 @@ class KotlinStandaloneDirectInheritorsProvider(private val project: Project) : K
         includeLocalInheritors: Boolean,
     ): Iterable<KtClassOrObject> {
         val classId = ktClass.getClassId() ?: return emptyList()
-        val possibleInheritors = staticDeclarationProviderFactory.getDirectInheritorCandidates(classId.shortClassName)
+
+        val aliases = mutableSetOf(classId.shortClassName)
+        calculateAliases(classId.shortClassName, aliases)
+
+        val possibleInheritors = aliases.flatMap { staticDeclarationProviderFactory.getDirectInheritorCandidates(it) }
 
         if (possibleInheritors.isEmpty()) {
             return emptyList()
@@ -65,6 +70,16 @@ class KotlinStandaloneDirectInheritorsProvider(private val project: Project) : K
         return possibleInheritors.filter { isValidInheritor(it, baseFirClass, scope, includeLocalInheritors) }
     }
 
+    private fun calculateAliases(aliasedName: Name, aliases: MutableSet<Name>) {
+        staticDeclarationProviderFactory.getInheritableTypeAliases(aliasedName).forEach { alias ->
+            val aliasName = alias.nameAsSafeName
+            val isNewAliasName = aliases.add(aliasName)
+            if (isNewAliasName) {
+                calculateAliases(aliasName, aliases)
+            }
+        }
+    }
+
     private fun isValidInheritor(
         candidate: KtClassOrObject,
         baseFirClass: FirClass,

analysis/analysis-api/testData/components/inheritorsProvider/sealedInheritors/sealedClassTypeAliasedFromDependency.txt

@@ -1,2 +1,8 @@
 /OneSealedChild
 class OneSealedChild : MySealedClass()
+
+/ThreeSealedChild
+class ThreeSealedChild : T2()
+
+/TwoSealedChild
+class TwoSealedChild : T1()

analysis/analysis-api/testData/components/inheritorsProvider/sealedInheritors/sealedClassTypeAliasedSameModule.txt

@@ -1,2 +1,8 @@
 /OneSealedChild
 class OneSealedChild : MySealedClass()
+
+/ThreeSealedChild
+class ThreeSealedChild : T2()
+
+/TwoSealedChild
+class TwoSealedChild : T1()

analysis/analysis-api/testData/components/inheritorsProvider/sealedInheritors/sealedInterfaceTypeAliasedFromDependency.txt

@@ -1,2 +1,8 @@
 /OneSealedChild
 class OneSealedChild : MySealedInterface
+
+/ThreeSealedChild
+class ThreeSealedChild : T2
+
+/TwoSealedChild
+class TwoSealedChild : T1

analysis/analysis-api/testData/components/inheritorsProvider/sealedInheritors/sealedInterfaceTypeAliasedSameModule.txt

@@ -1,2 +1,8 @@
 /OneSealedChild
 class OneSealedChild : MySealedInterface
+
+/ThreeSealedChild
+class ThreeSealedChild : T2
+
+/TwoSealedChild
+class TwoSealedChild : T1

compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/ktPsiUtil.kt

@@ -183,16 +183,7 @@ fun StubBasedPsiElementBase<out KotlinClassOrObjectStub<out KtClassOrObject>>.ge
 
         val file = containingFile
         if (file is KtFile) {
-            val directive = file.findImportByAlias(referencedName)
-            if (directive != null) {
-                var reference = directive.importedReference
-                while (reference is KtDotQualifiedExpression) {
-                    reference = reference.selectorExpression
-                }
-                if (reference is KtSimpleNameExpression) {
-                    result.add(reference.getReferencedName())
-                }
-            }
+            getImportedSimpleNameByImportAlias(file, referencedName)?.let(result::add)
         }
     }
 
@@ -721,4 +712,18 @@ internal fun isKtFile(parent: PsiElement?): Boolean {
     //avoid loading KtFile which depends on java psi, which is not available in some setup
     //e.g. remote dev https://youtrack.jetbrains.com/issue/GTW-7554
     return parent is PsiFile && parent.language == KotlinLanguage.INSTANCE
-}
+}
+
+fun getImportedSimpleNameByImportAlias(file: KtFile, aliasName: String): String? {
+    val directive = file.findImportByAlias(aliasName) ?: return null
+
+    var reference = directive.importedReference
+    while (reference is KtDotQualifiedExpression) {
+        reference = reference.selectorExpression
+    }
+    if (reference is KtSimpleNameExpression) {
+        return reference.getReferencedName()
+    }
+
+    return null
+}