1 files changed, 148 insertions, 0 deletions
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/java/KSTypeReferenceJavaImpl.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/java/KSTypeReferenceJavaImpl.kt
new file mode 100644
index 00000000..6b035244
--- /dev/null
+++ b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/symbol/impl/java/KSTypeReferenceJavaImpl.kt
@@ -0,0 +1,148 @@
+/*
+ * Copyright 2020 Google LLC
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.devtools.ksp.symbol.impl.java
+
+import com.google.devtools.ksp.ExceptionMessage
+import com.google.devtools.ksp.KSObjectCache
+import com.google.devtools.ksp.memoized
+import com.google.devtools.ksp.processing.impl.ResolverImpl
+import com.google.devtools.ksp.symbol.KSAnnotated
+import com.google.devtools.ksp.symbol.KSAnnotation
+import com.google.devtools.ksp.symbol.KSNode
+import com.google.devtools.ksp.symbol.KSReferenceElement
+import com.google.devtools.ksp.symbol.KSType
+import com.google.devtools.ksp.symbol.KSTypeReference
+import com.google.devtools.ksp.symbol.KSVisitor
+import com.google.devtools.ksp.symbol.Location
+import com.google.devtools.ksp.symbol.Modifier
+import com.google.devtools.ksp.symbol.NonExistLocation
+import com.google.devtools.ksp.symbol.Origin
+import com.google.devtools.ksp.symbol.impl.binary.KSClassDeclarationDescriptorImpl
+import com.google.devtools.ksp.symbol.impl.binary.KSClassifierReferenceDescriptorImpl
+import com.google.devtools.ksp.symbol.impl.kotlin.KSErrorType
+import com.google.devtools.ksp.symbol.impl.kotlin.KSTypeImpl
+import com.google.devtools.ksp.symbol.impl.toLocation
+import com.intellij.psi.PsiArrayType
+import com.intellij.psi.PsiClassType
+import com.intellij.psi.PsiPrimitiveType
+import com.intellij.psi.PsiType
+import com.intellij.psi.PsiWildcardType
+import com.intellij.psi.impl.source.PsiClassReferenceType
+import org.jetbrains.kotlin.descriptors.NotFoundClasses
+import org.jetbrains.kotlin.load.java.NOT_NULL_ANNOTATIONS
+import org.jetbrains.kotlin.load.java.NULLABLE_ANNOTATIONS
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.types.typeUtil.makeNullable
+
+class KSTypeReferenceJavaImpl private constructor(val psi: PsiType, override val parent: KSNode?) : KSTypeReference {
+    companion object : KSObjectCache<Pair<PsiType, KSNode?>, KSTypeReferenceJavaImpl>() {
+        fun getCached(psi: PsiType, parent: KSNode?) = cache
+            .getOrPut(Pair(psi, parent)) { KSTypeReferenceJavaImpl(psi, parent) }
+    }
+
+    override val origin = Origin.JAVA
+
+    override val location: Location by lazy {
+        (psi as? PsiClassReferenceType)?.reference?.toLocation() ?: NonExistLocation
+    }
+
+    override val annotations: Sequence<KSAnnotation> by lazy {
+        psi.annotations.asSequence().map { KSAnnotationJavaImpl.getCached(it) }.memoized()
+    }
+
+    override val modifiers: Set<Modifier> = emptySet()
+
+    override val element: KSReferenceElement by lazy {
+        fun PsiPrimitiveType.toKotlinType(): KotlinType {
+            return when (this.name) {
+                "int" -> ResolverImpl.instance!!.module.builtIns.intType
+                "short" -> ResolverImpl.instance!!.module.builtIns.shortType
+                "byte" -> ResolverImpl.instance!!.module.builtIns.byteType
+                "long" -> ResolverImpl.instance!!.module.builtIns.longType
+                "float" -> ResolverImpl.instance!!.module.builtIns.floatType
+                "double" -> ResolverImpl.instance!!.module.builtIns.doubleType
+                "char" -> ResolverImpl.instance!!.module.builtIns.charType
+                "boolean" -> ResolverImpl.instance!!.module.builtIns.booleanType
+                "void" -> ResolverImpl.instance!!.module.builtIns.unitType
+                else -> throw IllegalStateException("Unexpected primitive type ${this.name}, $ExceptionMessage")
+            }
+        }
+
+        val type = if (psi is PsiWildcardType) {
+            psi.bound
+        } else {
+            psi
+        }
+        when (type) {
+            is PsiClassType -> KSClassifierReferenceJavaImpl.getCached(type, this)
+            is PsiWildcardType -> KSClassifierReferenceJavaImpl.getCached(type.extendsBound as PsiClassType, this)
+            is PsiPrimitiveType -> KSClassifierReferenceDescriptorImpl.getCached(type.toKotlinType(), origin, this)
+            is PsiArrayType -> {
+                val componentType = ResolverImpl.instance!!.resolveJavaType(type.componentType, this)
+                if (type.componentType !is PsiPrimitiveType) {
+                    KSClassifierReferenceDescriptorImpl.getCached(
+                        ResolverImpl.instance!!.module.builtIns.getArrayType(Variance.INVARIANT, componentType),
+                        origin,
+                        this
+                    )
+                } else {
+                    KSClassifierReferenceDescriptorImpl.getCached(
+                        ResolverImpl.instance!!.module.builtIns
+                            .getPrimitiveArrayKotlinTypeByPrimitiveKotlinType(componentType)!!,
+                        origin, this
+                    )
+                }
+            }
+            null ->
+                KSClassifierReferenceDescriptorImpl.getCached(
+                    (ResolverImpl.instance!!.builtIns.anyType as KSTypeImpl).kotlinType.makeNullable(), origin, this
+                )
+            else -> throw IllegalStateException("Unexpected psi type for ${type.javaClass}, $ExceptionMessage")
+        }
+    }
+
+    override fun resolve(): KSType {
+        val resolvedType = ResolverImpl.instance!!.resolveUserType(this)
+        val relatedAnnotations = (annotations + ((parent as? KSAnnotated)?.annotations ?: emptySequence()))
+            .mapNotNull {
+                (it.annotationType.resolve() as? KSTypeImpl)?.kotlinType?.constructor?.declarationDescriptor?.fqNameSafe
+            }
+        val resolved = if ((resolvedType.declaration as? KSClassDeclarationDescriptorImpl)
+            ?.descriptor is NotFoundClasses.MockClassDescriptor
+        ) {
+            KSErrorType
+        } else resolvedType
+        val hasNotNull = relatedAnnotations.any { it in NOT_NULL_ANNOTATIONS }
+        val hasNullable = relatedAnnotations.any { it in NULLABLE_ANNOTATIONS }
+        return if (hasNullable && !hasNotNull) {
+            resolved.makeNullable()
+        } else if (!hasNullable && hasNotNull) {
+            resolved.makeNotNullable()
+        } else resolved
+    }
+
+    override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
+        return visitor.visitTypeReference(this, data)
+    }
+
+    override fun toString(): String {
+        return element.toString()
+    }
+}