Stubs: add classId to KtClassOrObject stubs

compiler/psi/src/org/jetbrains/kotlin/psi/KtClassLikeDeclaration.kt

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2010-2021 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.psi
+
+import org.jetbrains.kotlin.name.ClassId
+
+interface KtClassLikeDeclaration : KtDeclaration {
+    /**
+     * Return [ClassId], if the class is not local (E.e, if a class can be accessed by that [ClassId] from another context)
+     *
+     * For classes that itself local (are declared inside a function or other local scope), returns `null`.
+     * For nested classes in local classes returns `null`.
+     * For KtEnumEntry returns null as enum entry is not a class semantically. And so, for nested classes in enum entry, returns `null`.
+     * Otherwise, returns non-null [ClassId].
+     *
+     * For returned ClassId, the [ClassId.isLocal] is always `false`.
+     */
+    fun getClassId(): ClassId?
+}

compiler/psi/src/org/jetbrains/kotlin/psi/KtClassOrObject.kt

@@ -26,12 +26,15 @@ import com.intellij.psi.stubs.IStubElementType
 import com.intellij.psi.tree.TokenSet
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.psiUtil.ClassIdCalculator
 import org.jetbrains.kotlin.psi.stubs.KotlinClassOrObjectStub
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 
 abstract class KtClassOrObject :
     KtTypeParameterListOwnerStub<KotlinClassOrObjectStub<out KtClassOrObject>>, KtDeclarationContainer, KtNamedDeclaration,
-    KtPureClassOrObject {
+    KtPureClassOrObject, KtClassLikeDeclaration {
     constructor(node: ASTNode) : super(node)
     constructor(stub: KotlinClassOrObjectStub<out KtClassOrObject>, nodeType: IStubElementType<*, *>) : super(stub, nodeType)
 
@@ -89,6 +92,11 @@ abstract class KtClassOrObject :
 
     fun isTopLevel(): Boolean = stub?.isTopLevel() ?: (parent is KtFile)
 
+    override fun getClassId(): ClassId? {
+        stub?.let { return it.getClassId() }
+        return ClassIdCalculator.calculateClassId(this)
+    }
+
     override fun isLocal(): Boolean = stub?.isLocal() ?: KtPsiUtil.isLocal(this)
 
     override fun getDeclarations(): List<KtDeclaration> = getBody()?.declarations.orEmpty()
@@ -180,6 +188,7 @@ abstract class KtClassOrObject :
     }
 }
 
+
 fun KtClassOrObject.getOrCreateBody(): KtClassBody {
     getBody()?.let { return it }
 

compiler/psi/src/org/jetbrains/kotlin/psi/KtEnumEntry.java

@@ -22,6 +22,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiEnumConstant;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.name.ClassId;
 import org.jetbrains.kotlin.psi.stubs.KotlinClassStub;
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes;
 
@@ -51,6 +52,12 @@ public class KtEnumEntry extends KtClass {
         return !getSuperTypeListEntries().isEmpty();
     }
 
+    @Nullable
+    @Override
+    public ClassId getClassId() {
+        return null;
+    }
+
     @Nullable
     public KtInitializerList getInitializerList() {
         return getStubOrPsiChild(KtStubElementTypes.INITIALIZER_LIST);

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

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2010-2021 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.psi.psiUtil
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.SpecialNames
+import org.jetbrains.kotlin.psi.*
+
+internal object ClassIdCalculator {
+    fun calculateClassId(declaration: KtClassLikeDeclaration): ClassId? {
+        var ktFile: KtFile? = null
+        var element: PsiElement? = declaration
+        val containingClasses = mutableListOf<KtClassOrObject>()
+        while (element != null) {
+            when (element) {
+                is KtEnumEntry -> {
+                    return null
+                }
+                is KtClassOrObject -> {
+                    containingClasses += element
+                }
+                is KtFile -> {
+                    ktFile = element
+                    break
+                }
+                is KtDeclaration -> {
+                    return null
+                }
+            }
+            element = element.parent
+        }
+        if (ktFile == null) return null
+        val relativeClassName = FqName.fromSegments(
+            containingClasses.reversed().map { containingClass ->
+                containingClass.name ?: SpecialNames.NO_NAME_PROVIDED.asString()
+            }
+        )
+        return ClassId(ktFile.packageFqName, relativeClassName, /*local=*/false)
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/StubInterfaces.kt

@@ -21,6 +21,7 @@ import com.intellij.psi.stubs.NamedStub
 import com.intellij.psi.stubs.PsiFileStub
 import com.intellij.psi.stubs.StubElement
 import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 
@@ -40,11 +41,15 @@ interface KotlinStubWithFqName<T : PsiNamedElement> : NamedStub<T> {
     fun getFqName(): FqName?
 }
 
+interface KotlinClassifierStub {
+    fun getClassId(): ClassId?
+}
+
 interface KotlinTypeAliasStub : KotlinStubWithFqName<KtTypeAlias> {
     fun isTopLevel(): Boolean
 }
 
-interface KotlinClassOrObjectStub<T : KtClassOrObject> : KotlinStubWithFqName<T> {
+interface KotlinClassOrObjectStub<T : KtClassOrObject> : KotlinClassifierStub, KotlinStubWithFqName<T> {
     fun isLocal(): Boolean
     fun getSuperNames(): List<String>
     fun isTopLevel(): Boolean

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/StubUtils.kt

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010-2021 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.psi.stubs
+
+import com.intellij.psi.stubs.StubInputStream
+import com.intellij.psi.stubs.StubOutputStream
+import org.jetbrains.kotlin.name.ClassId
+
+object StubUtils {
+    @JvmStatic
+    fun deserializeClassId(dataStream: StubInputStream): ClassId? {
+        val classId = dataStream.readName() ?: return null
+        return ClassId.fromString(classId.string)
+    }
+
+    @JvmStatic
+    fun serializeClassId(dataStream: StubOutputStream, classId: ClassId?) {
+        dataStream.writeName(classId?.asString())
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/elements/KtClassElementType.java

@@ -24,11 +24,13 @@ import com.intellij.psi.stubs.StubOutputStream;
 import com.intellij.util.io.StringRef;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.name.ClassId;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.psi.KtClass;
 import org.jetbrains.kotlin.psi.KtEnumEntry;
 import org.jetbrains.kotlin.psi.psiUtil.KtPsiUtilKt;
 import org.jetbrains.kotlin.psi.stubs.KotlinClassStub;
+import org.jetbrains.kotlin.psi.stubs.StubUtils;
 import org.jetbrains.kotlin.psi.stubs.impl.KotlinClassStubImpl;
 import org.jetbrains.kotlin.psi.stubs.impl.Utils;
 
@@ -60,7 +62,7 @@ public class KtClassElementType extends KtStubElementType<KotlinClassStub, KtCla
         List<String> superNames = KtPsiUtilKt.getSuperNames(psi);
         return new KotlinClassStubImpl(
                 getStubType(isEnumEntry), (StubElement<?>) parentStub,
-                StringRef.fromString(fqName != null ? fqName.asString() : null),
+                StringRef.fromString(fqName != null ? fqName.asString() : null), psi.getClassId(),
                 StringRef.fromString(psi.getName()),
                 Utils.INSTANCE.wrapStrings(superNames),
                 psi.isInterface(), isEnumEntry, psi.isLocal(), psi.isTopLevel()
@@ -70,8 +72,12 @@ public class KtClassElementType extends KtStubElementType<KotlinClassStub, KtCla
     @Override
     public void serialize(@NotNull KotlinClassStub stub, @NotNull StubOutputStream dataStream) throws IOException {
         dataStream.writeName(stub.getName());
+
         FqName fqName = stub.getFqName();
         dataStream.writeName(fqName == null ? null : fqName.asString());
+
+        StubUtils.serializeClassId(dataStream, stub.getClassId());
+
         dataStream.writeBoolean(stub.isInterface());
         dataStream.writeBoolean(stub.isEnumEntry());
         dataStream.writeBoolean(stub.isLocal());
@@ -89,6 +95,9 @@ public class KtClassElementType extends KtStubElementType<KotlinClassStub, KtCla
     public KotlinClassStub deserialize(@NotNull StubInputStream dataStream, StubElement parentStub) throws IOException {
         StringRef name = dataStream.readName();
         StringRef qualifiedName = dataStream.readName();
+
+        ClassId classId = StubUtils.deserializeClassId(dataStream);
+
         boolean isTrait = dataStream.readBoolean();
         boolean isEnumEntry = dataStream.readBoolean();
         boolean isLocal = dataStream.readBoolean();
@@ -101,7 +110,7 @@ public class KtClassElementType extends KtStubElementType<KotlinClassStub, KtCla
         }
 
         return new KotlinClassStubImpl(
-                getStubType(isEnumEntry), (StubElement<?>) parentStub, qualifiedName, name, superNames,
+                getStubType(isEnumEntry), (StubElement<?>) parentStub, qualifiedName,classId, name, superNames,
                 isTrait, isEnumEntry, isLocal, isTopLevel
         );
     }

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/elements/KtObjectElementType.java

@@ -23,10 +23,12 @@ import com.intellij.psi.stubs.StubOutputStream;
 import com.intellij.util.io.StringRef;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.name.ClassId;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.psi.KtObjectDeclaration;
 import org.jetbrains.kotlin.psi.psiUtil.KtPsiUtilKt;
 import org.jetbrains.kotlin.psi.stubs.KotlinObjectStub;
+import org.jetbrains.kotlin.psi.stubs.StubUtils;
 import org.jetbrains.kotlin.psi.stubs.impl.KotlinObjectStubImpl;
 import org.jetbrains.kotlin.psi.stubs.impl.Utils;
 
@@ -45,7 +47,7 @@ public class KtObjectElementType extends KtStubElementType<KotlinObjectStub, KtO
         FqName fqName = KtPsiUtilKt.safeFqNameForLazyResolve(psi);
         List<String> superNames = KtPsiUtilKt.getSuperNames(psi);
         return new KotlinObjectStubImpl(
-                (StubElement<?>) parentStub, StringRef.fromString(name), fqName, Utils.INSTANCE.wrapStrings(superNames),
+                (StubElement<?>) parentStub, StringRef.fromString(name), fqName, psi.getClassId(), Utils.INSTANCE.wrapStrings(superNames),
                 psi.isTopLevel(), psi.isCompanion(), psi.isLocal(), psi.isObjectLiteral()
         );
     }
@@ -57,6 +59,8 @@ public class KtObjectElementType extends KtStubElementType<KotlinObjectStub, KtO
         FqName fqName = stub.getFqName();
         dataStream.writeName(fqName != null ? fqName.toString() : null);
 
+        StubUtils.serializeClassId(dataStream, stub.getClassId());
+
         dataStream.writeBoolean(stub.isTopLevel());
         dataStream.writeBoolean(stub.isCompanion());
         dataStream.writeBoolean(stub.isLocal());
@@ -73,9 +77,12 @@ public class KtObjectElementType extends KtStubElementType<KotlinObjectStub, KtO
     @Override
     public KotlinObjectStub deserialize(@NotNull StubInputStream dataStream, StubElement parentStub) throws IOException {
         StringRef name = dataStream.readName();
+
         StringRef fqNameStr = dataStream.readName();
         FqName fqName = fqNameStr != null ? new FqName(fqNameStr.toString()) : null;
 
+        ClassId classId = StubUtils.deserializeClassId(dataStream);
+
         boolean isTopLevel = dataStream.readBoolean();
         boolean isCompanion = dataStream.readBoolean();
         boolean isLocal = dataStream.readBoolean();
@@ -88,7 +95,7 @@ public class KtObjectElementType extends KtStubElementType<KotlinObjectStub, KtO
         }
 
         return new KotlinObjectStubImpl(
-                (StubElement<?>) parentStub, name, fqName, superNames, isTopLevel, isCompanion, isLocal, isObjectLiteral
+                (StubElement<?>) parentStub, name, fqName, classId, superNames, isTopLevel, isCompanion, isLocal, isObjectLiteral
         );
     }
 

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/impl/KotlinClassStubImpl.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.psi.stubs.impl
 import com.intellij.psi.PsiElement
 import com.intellij.psi.stubs.StubElement
 import com.intellij.util.io.StringRef
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.stubs.KotlinClassStub
@@ -29,12 +30,13 @@ class KotlinClassStubImpl(
     type: KtClassElementType,
     parent: StubElement<out PsiElement>?,
     private val qualifiedName: StringRef?,
+    private val classId: ClassId?,
     private val name: StringRef?,
     private val superNames: Array<StringRef>,
     private val isInterface: Boolean,
     private val isEnumEntry: Boolean,
     private val isLocal: Boolean,
-    private val isTopLevel: Boolean
+    private val isTopLevel: Boolean,
 ) : KotlinStubBaseImpl<KtClass>(parent, type), KotlinClassStub {
 
     override fun getFqName(): FqName? {
@@ -55,5 +57,7 @@ class KotlinClassStubImpl(
         return result
     }
 
+    override fun getClassId(): ClassId? = classId
+
     override fun isTopLevel() = isTopLevel
 }

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/impl/KotlinObjectStubImpl.kt

@@ -24,16 +24,18 @@ import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.name.FqName
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.name.ClassId
 
 class KotlinObjectStubImpl(
     parent: StubElement<out PsiElement>?,
     private val name: StringRef?,
     private val fqName: FqName?,
+    private val classId: ClassId?,
     private val superNames: Array<StringRef>,
     private val isTopLevel: Boolean,
     private val isDefault: Boolean,
     private val isLocal: Boolean,
-    private val isObjectLiteral: Boolean
+    private val isObjectLiteral: Boolean,
 ) : KotlinStubBaseImpl<KtObjectDeclaration>(parent, KtStubElementTypes.OBJECT_DECLARATION), KotlinObjectStub {
     override fun getFqName() = fqName
     override fun getName() = StringRef.toString(name)
@@ -42,4 +44,5 @@ class KotlinObjectStubImpl(
     override fun isCompanion() = isDefault
     override fun isObjectLiteral() = isObjectLiteral
     override fun isLocal() = isLocal
+    override fun getClassId(): ClassId? = classId
 }

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/impl/KotlinStubBaseImpl.kt

@@ -24,6 +24,7 @@ import com.intellij.psi.stubs.StubElement
 import org.jetbrains.kotlin.psi.KtElementImplStub
 import org.jetbrains.kotlin.psi.stubs.KotlinCallableStubBase
 import org.jetbrains.kotlin.psi.stubs.KotlinClassOrObjectStub
+import org.jetbrains.kotlin.psi.stubs.KotlinClassifierStub
 import org.jetbrains.kotlin.psi.stubs.KotlinStubWithFqName
 import java.lang.reflect.Method
 import java.util.ArrayList
@@ -82,6 +83,7 @@ open class KotlinStubBaseImpl<T : KtElementImplStub<*>>(parent: StubElement<*>?,
 
         private val BASE_STUB_INTERFACES = listOf(
             KotlinStubWithFqName::class.java,
+            KotlinClassifierStub::class.java,
             KotlinClassOrObjectStub::class.java,
             NamedStub::class.java,
             KotlinCallableStubBase::class.java