Adds KotlinPoet integration to Room Processing

room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/CompilationResultSubject.kt

@@ -427,7 +427,10 @@ internal class KotlinCompileTestingCompilationResult(
                         Source.loadJavaSource(sourceFile, qName)
                     }
                     sourceFile.name.endsWith(".kt") -> {
-                        Source.loadKotlinSource(sourceFile)
+                        val relativePath = sourceFile.absolutePath.substringAfter(
+                            srcRoot.absolutePath
+                        ).dropWhile { it == '/' }
+                        Source.loadKotlinSource(sourceFile, relativePath)
                     }
                     else -> null
                 }

room/compiler-processing-testing/src/main/java/androidx/room/compiler/processing/util/Source.kt

@@ -126,10 +126,11 @@ sealed class Source {
         }
 
         fun loadKotlinSource(
-            file: File
+            file: File,
+            relativePath: String
         ): Source {
             check(file.exists() && file.name.endsWith(".kt"))
-            return kotlin(file.absolutePath, file.readText())
+            return kotlin(relativePath, file.readText())
         }
 
         fun loadJavaSource(
@@ -142,13 +143,14 @@ sealed class Source {
 
         fun load(
             file: File,
-            qName: String
+            qName: String,
+            relativePath: String
         ): Source {
             check(file.exists()) {
                 "file does not exist: ${file.absolutePath}"
             }
             return when {
-                file.name.endsWith(".kt") -> loadKotlinSource(file)
+                file.name.endsWith(".kt") -> loadKotlinSource(file, relativePath)
                 file.name.endsWith(".java") -> loadJavaSource(file, qName)
                 else -> error("invalid file extension ${file.name}")
             }

room/compiler-processing-testing/src/test/java/androidx/room/compiler/processing/util/GeneratedCodeMatchTest.kt

@@ -20,7 +20,11 @@ import com.google.common.truth.Truth.assertThat
 import com.squareup.javapoet.JavaFile
 import com.squareup.javapoet.TypeName
 import com.squareup.javapoet.TypeSpec
+import com.squareup.kotlinpoet.BOOLEAN
+import com.squareup.kotlinpoet.FileSpec
+import com.squareup.kotlinpoet.TypeSpec as KTypeSpec
 import org.junit.Test
+import org.junit.AssumptionViolatedException
 import org.junit.runner.RunWith
 import org.junit.runners.Parameterized
 
@@ -111,4 +115,106 @@ class GeneratedCodeMatchTest internal constructor(
         )
         assertThat(result.exceptionOrNull()).hasMessageThat().contains(mismatch.toString())
     }
+
+    @Test
+    fun successfulGeneratedKotlinCodeMatch() {
+        // java environment will not generate kotlin files
+        if (runTest.toString() == "java") {
+            throw AssumptionViolatedException("javaAP won't generate kotlin code.")
+        }
+
+        val file = FileSpec.builder("foo.bar", "Baz")
+            .addType(KTypeSpec.classBuilder("Baz").build())
+            .build()
+        runTest { invocation ->
+            if (invocation.processingEnv.findTypeElement("foo.bar.Baz") == null) {
+                invocation.processingEnv.filer.write(file)
+            }
+            invocation.assertCompilationResult {
+                generatedSource(
+                    Source.kotlin("foo/bar/Baz.kt", file.toString())
+                )
+            }
+        }
+    }
+
+    @Test
+    fun missingGeneratedKotlinCode_mismatch() {
+        // java environment will not generate kotlin files
+        if (runTest.toString() == "java") {
+            throw AssumptionViolatedException("javaAP won't generate kotlin code.")
+        }
+
+        val generated = FileSpec.builder("foo.bar", "Baz")
+            .addType(
+                KTypeSpec.classBuilder("Baz")
+                    .addProperty("bar", BOOLEAN)
+                    .build()
+            )
+            .build()
+        val expected = FileSpec.builder("foo.bar", "Baz")
+            .addType(
+                KTypeSpec.classBuilder("Baz")
+                    .addProperty("foo", BOOLEAN)
+                    .build()
+            )
+            .build()
+
+        val result = runCatching {
+            runTest { invocation ->
+                if (invocation.processingEnv.findTypeElement("foo.bar.Baz") == null) {
+                    invocation.processingEnv.filer.write(generated)
+                }
+                invocation.assertCompilationResult {
+                    generatedSource(
+                        Source.kotlin("foo/bar/Baz.kt", expected.toString())
+                    )
+                }
+            }
+        }
+
+        val mismatch = SourceFileMismatch(
+            expected = Line(
+                pos = 6,
+                content = "val foo: Boolean"
+            ),
+            actual = Line(
+                pos = 6,
+                content = "val bar: Boolean"
+            )
+        )
+        assertThat(result.exceptionOrNull()).hasMessageThat().contains(mismatch.toString())
+    }
+
+    @Test
+    fun missingGeneratedKotlinCode_javaAP() {
+        if (runTest.toString() != "java") {
+            throw AssumptionViolatedException("Testing scenario for javaAP only.")
+        }
+
+        val file = FileSpec.builder("foo.bar", "Baz")
+            .addType(KTypeSpec.classBuilder("Baz").build())
+            .build()
+
+        val result = runCatching {
+            runTest { invocation ->
+                if (invocation.processingEnv.findTypeElement("foo.bar.Baz") == null) {
+                    invocation.processingEnv.filer.write(file)
+                }
+                invocation.assertCompilationResult {
+                    generatedSource(
+                        Source.kotlin("foo/bar/Baz.kt", file.toString())
+                    )
+                }
+            }
+        }
+
+        assertThat(result.exceptionOrNull())
+            .hasCauseThat()
+            .hasMessageThat()
+            .contains(
+                "Could not generate kotlin file foo/bar/Baz.kt. The annotation processing " +
+                    "environment is not set to generate Kotlin files."
+            )
+    }
 }

room/compiler-processing/build.gradle

@@ -28,6 +28,7 @@ plugins {
 dependencies {
     api(KOTLIN_STDLIB)
     api(JAVAPOET)
+    api(KOTLINPOET)
     implementation("androidx.annotation:annotation:1.1.0")
     implementation(GUAVA)
     implementation(AUTO_COMMON)

room/compiler-processing/src/main/java/androidx/room/compiler/processing/XFiler.kt

@@ -17,12 +17,17 @@
 package androidx.room.compiler.processing
 
 import com.squareup.javapoet.JavaFile
+import com.squareup.kotlinpoet.FileSpec
 
 /**
  * Code generation interface for XProcessing.
  */
 interface XFiler {
     fun write(javaFile: JavaFile)
+
+    fun write(fileSpec: FileSpec)
 }
 
 fun JavaFile.writeTo(generator: XFiler) = generator.write(this)
+
+fun FileSpec.writeTo(generator: XFiler) = generator.write(this)

room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacFiler.kt

@@ -18,10 +18,20 @@ package androidx.room.compiler.processing.javac
 
 import androidx.room.compiler.processing.XFiler
 import com.squareup.javapoet.JavaFile
-import javax.annotation.processing.Filer
+import com.squareup.kotlinpoet.FileSpec
+import javax.annotation.processing.ProcessingEnvironment
 
-internal class JavacFiler(val filer: Filer) : XFiler {
+internal class JavacFiler(val processingEnv: ProcessingEnvironment) : XFiler {
     override fun write(javaFile: JavaFile) {
-        javaFile.writeTo(filer)
+        javaFile.writeTo(processingEnv.filer)
+    }
+
+    override fun write(fileSpec: FileSpec) {
+        require(processingEnv.options.containsKey("kapt.kotlin.generated")) {
+            val filePath = fileSpec.packageName.replace('.', '/')
+            "Could not generate kotlin file $filePath/${fileSpec.name}.kt. The " +
+                "annotation processing environment is not set to generate Kotlin files."
+        }
+        fileSpec.writeTo(processingEnv.filer)
     }
 }

room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacProcessingEnv.kt

@@ -58,7 +58,7 @@ internal class JavacProcessingEnv(
         JavacProcessingEnvMessager(delegate)
     }
 
-    override val filer = JavacFiler(delegate.filer)
+    override val filer = JavacFiler(delegate)
 
     override val options: Map<String, String>
         get() = delegate.options

room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspFiler.kt

@@ -20,7 +20,12 @@ import androidx.room.compiler.processing.XFiler
 import androidx.room.compiler.processing.XMessager
 import com.google.devtools.ksp.processing.CodeGenerator
 import com.google.devtools.ksp.processing.Dependencies
+import com.google.devtools.ksp.symbol.KSFile
 import com.squareup.javapoet.JavaFile
+import com.squareup.kotlinpoet.FileSpec
+import com.squareup.kotlinpoet.TypeSpec
+import java.io.OutputStream
+import javax.lang.model.element.Element
 import javax.tools.Diagnostic
 
 internal class KspFiler(
@@ -29,17 +34,56 @@ internal class KspFiler(
 ) : XFiler {
     override fun write(javaFile: JavaFile) {
         val originatingFiles = javaFile.typeSpec.originatingElements
-            .map {
-                check(it is KSFileAsOriginatingElement) {
-                    "Unexpected element type in originating elements. $it"
-                }
-                it.ksFile
+            .map(::originatingFileFor)
+
+        createNewFile(
+            originatingFiles = originatingFiles,
+            packageName = javaFile.packageName,
+            fileName = javaFile.typeSpec.name,
+            extensionName = "java"
+        ).use { outputStream ->
+            outputStream.bufferedWriter(Charsets.UTF_8).use {
+                javaFile.writeTo(it)
+            }
+        }
+    }
+
+    override fun write(fileSpec: FileSpec) {
+        val originatingFiles = fileSpec.members
+            .filterIsInstance<TypeSpec>()
+            .flatMap { it.originatingElements }
+            .map(::originatingFileFor)
+
+        createNewFile(
+            originatingFiles = originatingFiles,
+            packageName = fileSpec.packageName,
+            fileName = fileSpec.name,
+            extensionName = "kt"
+        ).use { outputStream ->
+            outputStream.bufferedWriter(Charsets.UTF_8).use {
+                fileSpec.writeTo(it)
             }
+        }
+    }
+
+    private fun originatingFileFor(element: Element): KSFile {
+        check(element is KSFileAsOriginatingElement) {
+            "Unexpected element type in originating elements. $element"
+        }
+        return element.ksFile
+    }
+
+    private fun createNewFile(
+        originatingFiles: List<KSFile>,
+        packageName: String,
+        fileName: String,
+        extensionName: String
+    ): OutputStream {
         val dependencies = if (originatingFiles.isEmpty()) {
             messager.printMessage(
                 Diagnostic.Kind.WARNING,
                 """
-                    No dependencies are reported for ${javaFile.typeSpec.name} which will prevent
+                    No dependencies are reported for $fileName which will prevent
                     incremental compilation. Please file a bug at:
                     https://issuetracker.google.com/issues/new?component=413107
                 """.trimIndent()
@@ -52,15 +96,11 @@ internal class KspFiler(
             )
         }
 
-        delegate.createNewFile(
+        return delegate.createNewFile(
             dependencies = dependencies,
-            packageName = javaFile.packageName,
-            fileName = javaFile.typeSpec.name,
-            extensionName = "java"
-        ).use { outputStream ->
-            outputStream.bufferedWriter(Charsets.UTF_8).use {
-                javaFile.writeTo(it)
-            }
-        }
+            packageName = packageName,
+            fileName = fileName,
+            extensionName = extensionName
+        )
     }
 }

room/compiler/src/test/kotlin/androidx/room/testing/test_util.kt

@@ -259,7 +259,7 @@ fun loadJavaCode(fileName: String, qName: String): JavaFileObject {
 
 fun loadTestSource(fileName: String, qName: String): Source {
     val contents = File("src/test/data/$fileName")
-    return Source.load(contents, qName)
+    return Source.load(contents, qName, fileName)
 }
 
 fun createVerifierFromEntitiesAndViews(invocation: TestInvocation): DatabaseVerifier {