implementation classKind for KSClassDeclaration

* Implement classKind for kotlin source symbols and library symbols.
* Extract common logic for getting declarations for KtSymbolWithMember.

5 files changed, 51 insertions, 34 deletions

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt
index 49a07478..f3ebaefe 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSClassDeclarationImpl.kt
@@ -26,8 +26,16 @@ import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtObjectDeclaration
 
 class KSClassDeclarationImpl(private val ktNamedClassOrObjectSymbol: KtNamedClassOrObjectSymbol) : KSClassDeclaration {
-    override val classKind: ClassKind
-        get() = TODO("Not yet implemented")
+    override val classKind: ClassKind by lazy {
+        when (ktNamedClassOrObjectSymbol.classKind) {
+            KtClassKind.CLASS -> ClassKind.CLASS
+            KtClassKind.ENUM_CLASS -> ClassKind.ENUM_CLASS
+            KtClassKind.ANNOTATION_CLASS -> ClassKind.ANNOTATION_CLASS
+            KtClassKind.INTERFACE -> ClassKind.INTERFACE
+            KtClassKind.COMPANION_OBJECT, KtClassKind.ANONYMOUS_OBJECT, KtClassKind.OBJECT -> ClassKind.OBJECT
+            KtClassKind.ENUM_ENTRY -> ClassKind.ENUM_ENTRY
+        }
+    }
     override val primaryConstructor: KSFunctionDeclaration? by lazy {
         analyzeWithSymbolAsContext(ktNamedClassOrObjectSymbol) {
             ktNamedClassOrObjectSymbol.getMemberScope().getConstructors().singleOrNull { it.isPrimary }?.let {
@@ -76,7 +84,7 @@ class KSClassDeclarationImpl(private val ktNamedClassOrObjectSymbol: KtNamedClas
     }
 
     override val qualifiedName: KSName? by lazy {
-        (ktNamedClassOrObjectSymbol.psi as? KtClassOrObject)?.fqName?.asString()?.let { KSNameImpl(it) }
+        ktNamedClassOrObjectSymbol.classIdIfNonLocal?.asFqNameString()?.let { KSNameImpl(it) }
     }
 
     override val typeParameters: List<KSTypeParameter> by lazy {
@@ -133,15 +141,6 @@ class KSClassDeclarationImpl(private val ktNamedClassOrObjectSymbol: KtNamedClas
     }
 
     override val declarations: Sequence<KSDeclaration> by lazy {
-        analyzeWithSymbolAsContext(ktNamedClassOrObjectSymbol) {
-            ktNamedClassOrObjectSymbol.getDeclaredMemberScope().getAllSymbols().map {
-                when (it) {
-                    is KtNamedClassOrObjectSymbol -> KSClassDeclarationImpl(it)
-                    is KtFunctionSymbol -> KSFunctionDeclarationImpl(it)
-                    is KtPropertySymbol -> KSPropertyDeclarationImpl(it)
-                    else -> throw IllegalStateException()
-                }
-            }
-        }
+        ktNamedClassOrObjectSymbol.declarations()
     }
 }
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt
index 0769a722..e3421bd8 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFunctionDeclarationImpl.kt
@@ -19,15 +19,11 @@ package com.google.devtools.ksp.impl.symbol.kotlin
 
 import com.google.devtools.ksp.symbol.*
 import com.google.devtools.ksp.toKSModifiers
-import org.jetbrains.kotlin.analysis.api.InvalidWayOfUsingAnalysisSession
-import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
-import org.jetbrains.kotlin.analysis.api.KtAnalysisSessionProvider
 import org.jetbrains.kotlin.analysis.api.annotations.annotations
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtAnnotatedSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolKind
 import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
@@ -40,9 +36,11 @@ class KSFunctionDeclarationImpl(private val ktFunctionSymbol: KtFunctionLikeSymb
             else -> throw IllegalStateException("Unexpected symbol kind ${ktFunctionSymbol.symbolKind}")
         }
     }
+
     override val isAbstract: Boolean by lazy {
         (ktFunctionSymbol as? KtFunctionSymbol)?.modality == Modality.ABSTRACT
     }
+
     override val extensionReceiver: KSTypeReference? by lazy {
         analyzeWithSymbolAsContext(ktFunctionSymbol) {
             if (!ktFunctionSymbol.isExtension) {
@@ -52,6 +50,7 @@ class KSFunctionDeclarationImpl(private val ktFunctionSymbol: KtFunctionLikeSymb
             }
         }
     }
+
     override val returnType: KSTypeReference? by lazy {
         analyzeWithSymbolAsContext(ktFunctionSymbol) {
             KSTypeReferenceImpl(ktFunctionSymbol.returnType)
@@ -76,9 +75,11 @@ class KSFunctionDeclarationImpl(private val ktFunctionSymbol: KtFunctionLikeSymb
             KSNameImpl("<init>")
         }
     }
+
     override val qualifiedName: KSName? by lazy {
         (ktFunctionSymbol.psi as? KtFunction)?.fqName?.asString()?.let { KSNameImpl(it) }
     }
+
     override val typeParameters: List<KSTypeParameter> by lazy {
         (ktFunctionSymbol as? KtFunctionSymbol)?.typeParameters?.map { KSTypeParameterImpl(it) } ?: emptyList()
     }
@@ -88,7 +89,7 @@ class KSFunctionDeclarationImpl(private val ktFunctionSymbol: KtFunctionLikeSymb
     override val parentDeclaration: KSDeclaration?
         get() = TODO("Not yet implemented")
     override val containingFile: KSFile? by lazy {
-        (ktFunctionSymbol.psi?.containingFile as? KtFile)?.let { KSFileImpl(it) }
+        ktFunctionSymbol.toContainingFile()
     }
 
     override val docString: String? by lazy {
@@ -117,8 +118,10 @@ class KSFunctionDeclarationImpl(private val ktFunctionSymbol: KtFunctionLikeSymb
     override val annotations: Sequence<KSAnnotation> by lazy {
         (ktFunctionSymbol as KtAnnotatedSymbol).annotations.asSequence().map { KSAnnotationImpl(it) }
     }
+
     override val isActual: Boolean
         get() = TODO("Not yet implemented")
+
     override val isExpect: Boolean
         get() = TODO("Not yet implemented")
 
@@ -130,15 +133,6 @@ class KSFunctionDeclarationImpl(private val ktFunctionSymbol: KtFunctionLikeSymb
         TODO("Not yet implemented")
     }
 
-    override val declarations: Sequence<KSDeclaration>
-        get() = TODO("Not yet implemented")
-}
-
-@OptIn(InvalidWayOfUsingAnalysisSession::class)
-internal inline fun <R> analyzeWithSymbolAsContext(
-    contextSymbol: KtSymbol,
-    action: KtAnalysisSession.() -> R
-): R {
-    return KtAnalysisSessionProvider
-        .getInstance(contextSymbol.psi!!.project).analyzeWithSymbolAsContext(contextSymbol, action)
+    // TODO: Implement with PSI
+    override val declarations: Sequence<KSDeclaration> = emptySequence()
 }
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt
index 2c72b37e..0bff1404 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt
@@ -23,7 +23,6 @@ import com.google.devtools.ksp.toKSModifiers
 import org.jetbrains.kotlin.analysis.api.annotations.annotations
 import org.jetbrains.kotlin.analysis.api.symbols.KtNamedClassOrObjectSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtPropertySymbol
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
@@ -77,7 +76,7 @@ class KSPropertyDeclarationImpl(private val ktPropertySymbol: KtPropertySymbol)
         get() = TODO("Not yet implemented")
 
     override val containingFile: KSFile? by lazy {
-        (ktPropertySymbol.psi?.containingFile as? KtFile)?.let { KSFileImpl(it) }
+        ktPropertySymbol.toContainingFile()
     }
     override val docString: String? by lazy {
         ktPropertySymbol.psi?.getDocString()
@@ -98,7 +97,7 @@ class KSPropertyDeclarationImpl(private val ktPropertySymbol: KtPropertySymbol)
     override val parent: KSNode? by lazy {
         analyzeWithSymbolAsContext(ktPropertySymbol) {
             ktPropertySymbol.getContainingSymbol()?.let { KSClassDeclarationImpl(it as KtNamedClassOrObjectSymbol) }
-                ?: KSFileImpl(ktPropertySymbol.psi!!.containingFile as KtFile)
+                ?: ktPropertySymbol.toContainingFile()
         }
     }
 
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt
index 4c8b776c..fda40a7b 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt
@@ -30,8 +30,8 @@ import org.jetbrains.kotlin.analysis.api.annotations.annotations
 import org.jetbrains.kotlin.analysis.api.types.KtType
 
 class KSTypeReferenceImpl(private val ktType: KtType) : KSTypeReference {
-    override val element: KSReferenceElement?
-        get() = TODO("Not yet implemented")
+    // FIXME: return correct reference element.
+    override val element: KSReferenceElement? = null
 
     override fun resolve(): KSType {
         return KSTypeImpl(ktType)
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt
index d271569f..1c4826e9 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt
@@ -23,8 +23,15 @@ import com.google.devtools.ksp.symbol.*
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.analyseWithCustomToken
+import org.jetbrains.kotlin.analysis.api.annotations.annotations
+import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtFunctionLikeSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtNamedClassOrObjectSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtPropertySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
+import org.jetbrains.kotlin.analysis.api.symbols.markers.KtAnnotatedSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
 import org.jetbrains.kotlin.analysis.api.tokens.AlwaysAccessibleValidityTokenFactory
 import org.jetbrains.kotlin.psi.KtElement
 
@@ -70,3 +77,21 @@ internal inline fun <R> analyzeWithSymbolAsContext(
 ): R {
     return analyseWithCustomToken(contextSymbol.psi as KtElement, AlwaysAccessibleValidityTokenFactory, action)
 }
+
+internal fun KtSymbolWithMembers.declarations(): Sequence<KSDeclaration> {
+    return analyzeWithSymbolAsContext(this) {
+        this@declarations.getDeclaredMemberScope().getAllSymbols().map {
+            when (it) {
+                is KtNamedClassOrObjectSymbol -> KSClassDeclarationImpl(it)
+                is KtFunctionLikeSymbol -> KSFunctionDeclarationImpl(it)
+                is KtPropertySymbol -> KSPropertyDeclarationImpl(it)
+                is KtEnumEntrySymbol -> KSClassDeclarationEnumEntryImpl(it)
+                else -> throw IllegalStateException()
+            }
+        }
+    }
+}
+
+internal fun KtAnnotatedSymbol.annotations(): Sequence<KSAnnotation> {
+    return this.annotations.asSequence().map { KSAnnotationImpl(it) }
+}