Implement XArrayType in KSP

This CL implements XArrayType in KSP.

In KSP, arrays are just generics so we detect when a type
is wrapped into KspType, similar to what we do in java
processor.

For TypeName, I opted in to java style array rather than
a Parameterized type. What we will do for TypeNames is
up for debate (e.g. at what level we'll swap them with java
representations). It is inconsistent with the rest of KSP
but for arrays, conversion is always so made sense to me
to do it here.

Bug: 160322705
Test: XArrayTypeTest
Change-Id: I2d687d5d5a77d460f66914f7167375c05036fe0d
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSAsMemberOf.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSAsMemberOf.kt
index 8d60651..2ea9db7 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSAsMemberOf.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSAsMemberOf.kt
@@ -26,7 +26,6 @@
 import org.jetbrains.kotlin.ksp.symbol.KSType
 import org.jetbrains.kotlin.ksp.symbol.KSTypeArgument
 import org.jetbrains.kotlin.ksp.symbol.KSTypeParameter
-import org.jetbrains.kotlin.ksp.symbol.KSTypeReference
 import org.jetbrains.kotlin.ksp.symbol.Nullability
 
 /**
@@ -124,17 +123,3 @@
     }
     return resolver.getTypeArgument(myType.swapResolvedType(resolved.makeNullable()), variance)
 }
-
-private fun KSTypeReference.swapResolvedType(replacement: KSType): KSTypeReference {
-    return DelegatingTypeReference(
-        original = this,
-        resolved = replacement
-    )
-}
-
-private class DelegatingTypeReference(
-    val original: KSTypeReference,
-    val resolved: KSType
-) : KSTypeReference by original {
-    override fun resolve(): KSType? = resolved
-}
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSTypeReferenceExt.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSTypeReferenceExt.kt
new file mode 100644
index 0000000..7c0612b
--- /dev/null
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSTypeReferenceExt.kt
@@ -0,0 +1,73 @@
+/*
+ * Copyright 2020 The Android Open Source Project
+ *
+ * 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 androidx.room.compiler.processing.ksp
+
+import org.jetbrains.kotlin.ksp.symbol.KSAnnotation
+import org.jetbrains.kotlin.ksp.symbol.KSReferenceElement
+import org.jetbrains.kotlin.ksp.symbol.KSType
+import org.jetbrains.kotlin.ksp.symbol.KSTypeReference
+import org.jetbrains.kotlin.ksp.symbol.KSVisitor
+import org.jetbrains.kotlin.ksp.symbol.Location
+import org.jetbrains.kotlin.ksp.symbol.Modifier
+import org.jetbrains.kotlin.ksp.symbol.NonExistLocation
+import org.jetbrains.kotlin.ksp.symbol.Origin
+
+/**
+ * Creates a new TypeReference from [this] where the resolved type [replacement] but everything
+ * else is the same (e.g. location).
+ */
+internal fun KSTypeReference.swapResolvedType(replacement: KSType): KSTypeReference {
+    return DelegatingTypeReference(
+        original = this,
+        resolved = replacement
+    )
+}
+
+/**
+ * Creates a [NonExistLocation] type reference for [this].
+ */
+internal fun KSType.createTypeReference(): KSTypeReference {
+    return NoLocationTypeReference(this)
+}
+
+private class DelegatingTypeReference(
+    val original: KSTypeReference,
+    val resolved: KSType
+) : KSTypeReference by original {
+    override fun resolve(): KSType? = resolved
+}
+
+private class NoLocationTypeReference(
+    val resolved: KSType
+) : KSTypeReference {
+    override val annotations: List<KSAnnotation>
+        get() = emptyList()
+    override val element: KSReferenceElement?
+        get() = null
+    override val location: Location
+        get() = NonExistLocation
+    override val modifiers: Set<Modifier>
+        get() = emptySet()
+    override val origin: Origin
+        get() = Origin.SYNTHETIC
+
+    override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
+        return visitor.visitTypeReference(this, data)
+    }
+
+    override fun resolve(): KSType = resolved
+}
\ No newline at end of file
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspArrayType.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspArrayType.kt
new file mode 100644
index 0000000..d578c20
--- /dev/null
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspArrayType.kt
@@ -0,0 +1,38 @@
+/*
+ * Copyright 2020 The Android Open Source Project
+ *
+ * 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 androidx.room.compiler.processing.ksp
+
+import androidx.room.compiler.processing.XArrayType
+import androidx.room.compiler.processing.XType
+import com.squareup.javapoet.ArrayTypeName
+import com.squareup.javapoet.TypeName
+import org.jetbrains.kotlin.ksp.symbol.KSType
+
+internal class KspArrayType(
+    env: KspProcessingEnv,
+    ksType: KSType
+) : KspType(
+    env, ksType
+), XArrayType {
+    override val componentType: XType by lazy {
+        typeArguments.first()
+    }
+
+    override val typeName: TypeName by lazy {
+        ArrayTypeName.of(typeArguments.first().typeName)
+    }
+}
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspProcessingEnv.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspProcessingEnv.kt
index 9c2bcb9..24ec0b3 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspProcessingEnv.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspProcessingEnv.kt
@@ -29,6 +29,7 @@
 import org.jetbrains.kotlin.ksp.symbol.KSClassDeclaration
 import org.jetbrains.kotlin.ksp.symbol.KSType
 import org.jetbrains.kotlin.ksp.symbol.KSTypeReference
+import org.jetbrains.kotlin.ksp.symbol.Variance
 import javax.annotation.processing.Filer
 
 internal class KspProcessingEnv(
@@ -76,19 +77,36 @@
     }
 
     override fun getArrayType(type: XType): XArrayType {
-        TODO("Not yet implemented")
+        check(type is KspType)
+        val arrayType = resolver.requireClass(KOTLIN_ARRAY_Q_NAME)
+        val ksType = arrayType.asType(
+            listOf(resolver.getTypeArgument(
+                type.ksType.createTypeReference(),
+                Variance.INVARIANT
+            ))
+        )
+        return KspArrayType(
+            env = this,
+            ksType = ksType
+        )
     }
 
     fun wrap(ksType: KSType): KspType {
-        return KspType(
-            env = this,
-            ksType = ksType)
+        return if (ksType.declaration.qualifiedName?.asString() == KOTLIN_ARRAY_Q_NAME) {
+            KspArrayType(
+                env = this,
+                ksType = ksType
+            )
+        } else {
+            KspType(
+                env = this,
+                ksType = ksType
+            )
+        }
     }
 
     fun wrap(ksTypeReference: KSTypeReference): KspType {
-        return KspType(
-            env = this,
-            ksTypeReference = ksTypeReference)
+        return wrap(ksTypeReference.requireType())
     }
 
     fun wrapClassDeclaration(declaration: KSClassDeclaration): KspTypeElement {
@@ -109,4 +127,8 @@
             resolver.builtIns.byteType.makeNullable()
         }
     }
+
+    companion object {
+        private const val KOTLIN_ARRAY_Q_NAME = "kotlin.Array"
+    }
 }
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspType.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspType.kt
index c83c021..85ab73d 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspType.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspType.kt
@@ -36,15 +36,10 @@
  * We don't necessarily have a [KSTypeReference] (e.g. if we are getting it from an element).
  * Similarly, we may not be able to get a [KSType] (e.g. if it resolves to error).
  */
-internal class KspType(
+internal open class KspType(
     private val env: KspProcessingEnv,
     val ksType: KSType
 ) : XDeclaredType, XEquality {
-    constructor(env: KspProcessingEnv, ksTypeReference: KSTypeReference) : this(
-        env = env,
-        ksType = ksTypeReference.requireType()
-    )
-
     override val rawType by lazy {
         KspRawType(this)
     }
diff --git a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/XArrayTypeTest.kt b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/XArrayTypeTest.kt
index 13a1dc4..8dec922 100644
--- a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/XArrayTypeTest.kt
+++ b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/XArrayTypeTest.kt
@@ -19,6 +19,7 @@
 import androidx.room.compiler.processing.util.Source
 import androidx.room.compiler.processing.util.getField
 import androidx.room.compiler.processing.util.runProcessorTest
+import androidx.room.compiler.processing.util.runProcessorTestIncludingKsp
 import com.google.common.truth.Truth.assertThat
 import com.squareup.javapoet.ArrayTypeName
 import com.squareup.javapoet.TypeName
@@ -26,7 +27,7 @@
 
 class XArrayTypeTest {
     @Test
-    fun xArrayType() {
+    fun java() {
         val source = Source.java(
             "foo.bar.Baz", """
             package foo.bar;
@@ -35,21 +36,27 @@
             }
         """.trimIndent()
         )
-        runProcessorTest(
+        runProcessorTestIncludingKsp(
             sources = listOf(source)
-        ) {
-            val type = it.processingEnv
+        ) { invocation ->
+            val type = invocation.processingEnv
                 .requireTypeElement("foo.bar.Baz")
                 .getField("param")
                 .type
             assertThat(type.isArray()).isTrue()
             assertThat(type.typeName).isEqualTo(
-                ArrayTypeName.of(TypeName.get(String::class.java))
+                ArrayTypeName.of(invocation.types.string)
             )
-            assertThat(type.asArray().componentType.typeName).isEqualTo(
-                TypeName.get(String::class.java)
-            )
+            type.asArray().componentType.let { component ->
+                assertThat(component.typeName).isEqualTo(invocation.types.string)
+                assertThat(component.nullability).isEqualTo(XNullability.UNKNOWN)
+            }
+        }
+    }
 
+    @Test
+    fun synthetic() {
+        runProcessorTestIncludingKsp {
             val objArray = it.processingEnv.getArrayType(
                 TypeName.OBJECT
             )
@@ -64,6 +71,41 @@
     }
 
     @Test
+    fun kotlin() {
+        val source = Source.kotlin(
+            "Foo.kt", """
+            package foo.bar
+            class Baz {
+                val nonNull:Array<String> = TODO()
+                val nullable:Array<String?> = TODO()
+            }
+        """.trimIndent()
+        )
+        runProcessorTestIncludingKsp(
+            sources = listOf(source)
+        ) { invocation ->
+            val element = invocation.processingEnv.requireTypeElement("foo.bar.Baz")
+            val nonNull = element.getField("nonNull").type
+            val nullable = element.getField("nullable").type
+            listOf(nonNull, nullable).forEach {
+                assertThat(nonNull.isArray()).isTrue()
+                assertThat(nonNull.typeName).isEqualTo(
+                    ArrayTypeName.of(invocation.types.string)
+                )
+            }
+
+            nonNull.asArray().componentType.let { component ->
+                assertThat(component.typeName).isEqualTo(invocation.types.string)
+                assertThat(component.nullability).isEqualTo(XNullability.NONNULL)
+            }
+            nullable.asArray().componentType.let { component ->
+                assertThat(component.typeName).isEqualTo(invocation.types.string)
+                assertThat(component.nullability).isEqualTo(XNullability.NULLABLE)
+            }
+        }
+    }
+
+    @Test
     fun notAnArray() {
         runProcessorTest {
             val list = it.processingEnv.requireType("java.util.List")
diff --git a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/KotlinTypeNames.kt b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/KotlinTypeNames.kt
index a1b0f22..32d77be 100644
--- a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/KotlinTypeNames.kt
+++ b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/KotlinTypeNames.kt
@@ -19,6 +19,8 @@
 import com.squareup.javapoet.ClassName
 
 internal object KotlinTypeNames {
+    val ANY_CLASS_NAME = ClassName.get("kotlin", "Any")
+    val UNIT_CLASS_NAME = ClassName.get("kotlin", "Unit")
     val INT_CLASS_NAME = ClassName.get("kotlin", "Int")
     val STRING_CLASS_NAME = ClassName.get("kotlin", "String")
     val LIST_CLASS_NAME = ClassName.get("kotlin.collections", "List")
diff --git a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/ProcessorTestExt.kt b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/ProcessorTestExt.kt
index 4cff818..57848c9 100644
--- a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/ProcessorTestExt.kt
+++ b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/ProcessorTestExt.kt
@@ -118,6 +118,20 @@
     runKaptTest(sources = sources, handler = handler, succeed = true)
 }
 
+/**
+ * This method is oddly named instead of being an overload on runProcessorTest to easily track
+ * which tests started to support KSP.
+ *
+ * Eventually, it will be merged with runProcessorTest when all tests pass with KSP.
+ */
+fun runProcessorTestIncludingKsp(
+    sources: List<Source> = emptyList(),
+    handler: (TestInvocation) -> Unit
+) {
+    runProcessorTest(sources = sources, handler = handler)
+    runKspTest(sources = sources, succeed = true, handler = handler)
+}
+
 fun runProcessorTestForFailedCompilation(
     sources: List<Source>,
     handler: (TestInvocation) -> Unit
diff --git a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/TestInvocation.kt b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/TestInvocation.kt
index e3ab7d4..9398adf 100644
--- a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/TestInvocation.kt
+++ b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/util/TestInvocation.kt
@@ -18,6 +18,8 @@
 
 import androidx.room.compiler.processing.XProcessingEnv
 import androidx.room.compiler.processing.ksp.KspProcessingEnv
+import com.squareup.javapoet.ClassName
+import com.squareup.javapoet.TypeName
 import org.jetbrains.kotlin.ksp.processing.Resolver
 
 class TestInvocation(
@@ -25,4 +27,36 @@
 ) {
     val kspResolver: Resolver
         get() = (processingEnv as KspProcessingEnv).resolver
+
+    val types by lazy {
+        if (processingEnv is KspProcessingEnv) {
+            Types(
+                string = KotlinTypeNames.STRING_CLASS_NAME,
+                voidOrUnit = KotlinTypeNames.UNIT_CLASS_NAME,
+                objectOrAny = KotlinTypeNames.ANY_CLASS_NAME,
+                boxedInt = KotlinTypeNames.INT_CLASS_NAME,
+                int = KotlinTypeNames.INT_CLASS_NAME
+            )
+        } else {
+            Types(
+                string = ClassName.get("java.lang", "String"),
+                voidOrUnit = TypeName.VOID,
+                objectOrAny = TypeName.OBJECT,
+                boxedInt = TypeName.INT.box(),
+                int = TypeName.INT
+            )
+        }
+    }
+
+    /**
+     * Helper class to hold types that change between KSP and Javap.
+     * e.g. Kotlin.String vs java.lang.String
+     */
+    class Types(
+        val string: ClassName,
+        val voidOrUnit: TypeName,
+        val objectOrAny: ClassName,
+        val boxedInt: TypeName,
+        val int: TypeName
+    )
 }