Merge "[GH] [Room][Compiler Processing] Originating element improvements" into androidx-main
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/JavaPoetExt.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/JavaPoetExt.kt
index f4860f2..8824318 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/JavaPoetExt.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/JavaPoetExt.kt
@@ -15,9 +15,7 @@
  */
 package androidx.room.compiler.processing
 
-import androidx.room.compiler.processing.javac.JavacElement
 import androidx.room.compiler.processing.javac.JavacExecutableElement
-import androidx.room.compiler.processing.ksp.KspElement
 import androidx.room.compiler.processing.ksp.KspMethodElement
 import androidx.room.compiler.processing.ksp.KspMethodType
 import com.squareup.javapoet.ClassName
@@ -51,12 +49,9 @@
  * Adds the given element as the originating element for compilation.
  * see [TypeSpec.Builder.addOriginatingElement].
  */
-fun TypeSpec.Builder.addOriginatingElement(element: XElement) {
-    if (element is JavacElement) {
-        this.addOriginatingElement(element.element)
-    } else if (element is KspElement) {
-        element.containingFileAsOriginatingElement()?.let(this::addOriginatingElement)
-    }
+fun TypeSpec.Builder.addOriginatingElement(element: XElement): TypeSpec.Builder {
+    element.originatingElementForPoet()?.let(this::addOriginatingElement)
+    return this
 }
 
 internal fun TypeName.rawTypeName(): TypeName {
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/KotlinPoetExt.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/KotlinPoetExt.kt
new file mode 100644
index 0000000..02affec
--- /dev/null
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/KotlinPoetExt.kt
@@ -0,0 +1,30 @@
+/*
+ * Copyright 2021 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
+
+import com.squareup.kotlinpoet.OriginatingElementsHolder
+
+/**
+ * Adds the given element as an originating element for compilation.
+ * see [OriginatingElementsHolder.Builder.addOriginatingElement].
+ */
+fun <T : OriginatingElementsHolder.Builder<T>> T.addOriginatingElement(
+    element: XElement
+): T {
+    element.originatingElementForPoet()?.let(this::addOriginatingElement)
+    return this
+}
\ No newline at end of file
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XElement.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XElement.kt
index a94f643..b06d8c9 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XElement.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XElement.kt
@@ -16,6 +16,10 @@
 
 package androidx.room.compiler.processing
 
+import androidx.room.compiler.processing.javac.JavacElement
+import androidx.room.compiler.processing.ksp.KSFileAsOriginatingElement
+import androidx.room.compiler.processing.ksp.KspElement
+import javax.lang.model.element.Element
 import kotlin.contracts.contract
 
 /**
@@ -75,3 +79,19 @@
     }
     return this is XConstructorElement
 }
+
+/**
+ * Attempts to get a Javac [Element] representing the originating element for attribution
+ * when writing a file for incremental processing.
+ *
+ * In KSP a [KSFileAsOriginatingElement] will be returned, which is a synthetic javac element
+ * that allows us to pass originating elements to JavaPoet and KotlinPoet, and later extract
+ * the KSP file when writing with [XFiler].
+ */
+internal fun XElement.originatingElementForPoet(): Element? {
+    return when (this) {
+        is JavacElement -> element
+        is KspElement -> containingFileAsOriginatingElement()
+        else -> error("Originating element is not implemented for ${this.javaClass}")
+    }
+}
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XFiler.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XFiler.kt
index 336814e..79db763 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XFiler.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XFiler.kt
@@ -23,11 +23,27 @@
  * Code generation interface for XProcessing.
  */
 interface XFiler {
-    fun write(javaFile: JavaFile)
 
-    fun write(fileSpec: FileSpec)
+    fun write(javaFile: JavaFile, mode: Mode = Mode.Isolating)
+
+    fun write(fileSpec: FileSpec, mode: Mode = Mode.Isolating)
+
+    /**
+     * Specifies whether a file represents aggregating or isolating inputs for incremental
+     * build purposes. This does not apply in Javac processing because aggregating vs isolating
+     * is set on the processor level. For more on KSP's definitions of isolating vs aggregating
+     * see the documentation at
+     * https://github.com/google/ksp/blob/master/docs/incremental.md
+     */
+    enum class Mode {
+        Aggregating, Isolating
+    }
 }
 
-fun JavaFile.writeTo(generator: XFiler) = generator.write(this)
+fun JavaFile.writeTo(generator: XFiler, mode: XFiler.Mode = XFiler.Mode.Isolating) {
+    generator.write(this, mode)
+}
 
-fun FileSpec.writeTo(generator: XFiler) = generator.write(this)
+fun FileSpec.writeTo(generator: XFiler, mode: XFiler.Mode = XFiler.Mode.Isolating) {
+    generator.write(this, mode)
+}
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacFiler.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacFiler.kt
index d84c8a6..484ed66 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacFiler.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacFiler.kt
@@ -22,11 +22,13 @@
 import javax.annotation.processing.ProcessingEnvironment
 
 internal class JavacFiler(val processingEnv: ProcessingEnvironment) : XFiler {
-    override fun write(javaFile: JavaFile) {
+
+    // "mode" is ignored in javac, and only applicable in KSP
+    override fun write(javaFile: JavaFile, mode: XFiler.Mode) {
         javaFile.writeTo(processingEnv.filer)
     }
 
-    override fun write(fileSpec: FileSpec) {
+    override fun write(fileSpec: FileSpec, mode: XFiler.Mode) {
         require(processingEnv.options.containsKey("kapt.kotlin.generated")) {
             val filePath = fileSpec.packageName.replace('.', '/')
             "Could not generate kotlin file $filePath/${fileSpec.name}.kt. The " +
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSFileAsOriginatingElement.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSFileAsOriginatingElement.kt
index a6fa6e4..ec05ba1 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSFileAsOriginatingElement.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KSFileAsOriginatingElement.kt
@@ -30,7 +30,7 @@
  * This wrapper class helps us wrap a KSFile as an originating element and KspFiler unwraps it to
  * get the actual KSFile out of it.
  */
-internal class KSFileAsOriginatingElement(
+internal data class KSFileAsOriginatingElement(
     val ksFile: KSFile
 ) : Element {
     override fun getAnnotationMirrors(): List<AnnotationMirror> {
@@ -76,10 +76,6 @@
         return null
     }
 
-    override fun toString(): String {
-        return ksFile.toString()
-    }
-
     private class NameImpl(private val str: String) : Name, CharSequence by str {
         override fun contentEquals(cs: CharSequence): Boolean {
             return str == cs.toString()
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspFiler.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspFiler.kt
index 0b92034..9098435 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspFiler.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspFiler.kt
@@ -32,7 +32,7 @@
     private val delegate: CodeGenerator,
     private val messager: XMessager,
 ) : XFiler {
-    override fun write(javaFile: JavaFile) {
+    override fun write(javaFile: JavaFile, mode: XFiler.Mode) {
         val originatingFiles = javaFile.typeSpec.originatingElements
             .map(::originatingFileFor)
 
@@ -40,7 +40,8 @@
             originatingFiles = originatingFiles,
             packageName = javaFile.packageName,
             fileName = javaFile.typeSpec.name,
-            extensionName = "java"
+            extensionName = "java",
+            aggregating = mode == XFiler.Mode.Aggregating
         ).use { outputStream ->
             outputStream.bufferedWriter(Charsets.UTF_8).use {
                 javaFile.writeTo(it)
@@ -48,7 +49,7 @@
         }
     }
 
-    override fun write(fileSpec: FileSpec) {
+    override fun write(fileSpec: FileSpec, mode: XFiler.Mode) {
         val originatingFiles = fileSpec.members
             .filterIsInstance<TypeSpec>()
             .flatMap { it.originatingElements }
@@ -58,7 +59,8 @@
             originatingFiles = originatingFiles,
             packageName = fileSpec.packageName,
             fileName = fileSpec.name,
-            extensionName = "kt"
+            extensionName = "kt",
+            aggregating = mode == XFiler.Mode.Aggregating
         ).use { outputStream ->
             outputStream.bufferedWriter(Charsets.UTF_8).use {
                 fileSpec.writeTo(it)
@@ -77,7 +79,8 @@
         originatingFiles: List<KSFile>,
         packageName: String,
         fileName: String,
-        extensionName: String
+        extensionName: String,
+        aggregating: Boolean
     ): OutputStream {
         val dependencies = if (originatingFiles.isEmpty()) {
             messager.printMessage(
@@ -91,7 +94,7 @@
             Dependencies.ALL_FILES
         } else {
             Dependencies(
-                aggregating = false,
+                aggregating = aggregating,
                 sources = originatingFiles.distinct().toTypedArray()
             )
         }
diff --git a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/OriginatingElementsTest.kt b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/OriginatingElementsTest.kt
new file mode 100644
index 0000000..fcc43f0
--- /dev/null
+++ b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/OriginatingElementsTest.kt
@@ -0,0 +1,107 @@
+/*
+ * Copyright 2021 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
+
+import androidx.room.compiler.processing.ksp.KSFileAsOriginatingElement
+import androidx.room.compiler.processing.ksp.KspTypeElement
+import androidx.room.compiler.processing.util.Source
+import androidx.room.compiler.processing.util.runProcessorTest
+import com.google.common.truth.Truth.assertThat
+import com.squareup.javapoet.TypeSpec
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import javax.lang.model.element.TypeElement
+
+@RunWith(JUnit4::class)
+class OriginatingElementsTest {
+
+    @Test
+    fun xElementIsConvertedToOriginatingElement() {
+        runProcessorTest(
+            sources = listOf(
+                Source.java(
+                    "foo.bar.Baz",
+                    """
+                package foo.bar;
+                public class Baz {
+                    private void foo() {}
+                    public String bar(String[] param1) {
+                        return "";
+                    }
+                }
+                    """.trimIndent()
+                )
+            )
+        ) {
+            val element = it.processingEnv.requireTypeElement("foo.bar.Baz")
+
+            val originatingElement = element.originatingElementForPoet()
+            assertThat(originatingElement).isNotNull()
+
+            if (it.isKsp) {
+                assertThat(originatingElement).isInstanceOf(KSFileAsOriginatingElement::class.java)
+
+                val originatingFile = (originatingElement as KSFileAsOriginatingElement).ksFile
+                assertThat(originatingFile)
+                    .isEqualTo(
+                        (element as KspTypeElement).declaration.containingFile
+                    )
+            } else {
+                assertThat(originatingElement).isInstanceOf(TypeElement::class.java)
+                assertThat((originatingElement as TypeElement).qualifiedName.toString())
+                    .isEqualTo("foo.bar.Baz")
+            }
+        }
+    }
+
+    @Test
+    fun originatingElementIsAddedToPoet() {
+        runProcessorTest(
+            sources = listOf(
+                Source.java(
+                    "foo.bar.Baz",
+                    """
+                package foo.bar;
+                public class Baz {
+                    private void foo() {}
+                    public String bar(String[] param1) {
+                        return "";
+                    }
+                }
+                    """.trimIndent()
+                )
+            )
+        ) {
+            val xTypeElement = it.processingEnv.requireTypeElement("foo.bar.Baz")
+
+            val javaPoetTypeSpec = TypeSpec.classBuilder("Foo")
+                .addOriginatingElement(xTypeElement)
+                .build()
+
+            val kotlinPoetTypeSpec = com.squareup.kotlinpoet.TypeSpec.classBuilder("Foo")
+                .addOriginatingElement(xTypeElement)
+                .build()
+
+            assertThat(javaPoetTypeSpec.originatingElements).apply {
+                hasSize(1)
+                containsExactlyElementsIn(kotlinPoetTypeSpec.originatingElements)
+                containsExactly(xTypeElement.originatingElementForPoet())
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/room/compiler/src/test/kotlin/androidx/room/writer/AutoMigrationWriterTest.kt b/room/compiler/src/test/kotlin/androidx/room/writer/AutoMigrationWriterTest.kt
index f1d962b..4da432f 100644
--- a/room/compiler/src/test/kotlin/androidx/room/writer/AutoMigrationWriterTest.kt
+++ b/room/compiler/src/test/kotlin/androidx/room/writer/AutoMigrationWriterTest.kt
@@ -16,7 +16,6 @@
 
 package androidx.room.writer
 
-import androidx.room.compiler.processing.XElement
 import androidx.room.compiler.processing.util.Source
 import androidx.room.compiler.processing.util.runProcessorTest
 import androidx.room.migration.bundle.FieldBundle
@@ -26,7 +25,6 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
-import org.mockito.Mockito.mock
 
 @RunWith(JUnit4::class)
 class AutoMigrationWriterTest {
@@ -73,7 +71,10 @@
                     removedOrRenamedTables = listOf()
                 ),
             )
-            AutoMigrationWriter(mock(XElement::class.java), autoMigrationResultWithNewAddedColumn)
+            AutoMigrationWriter(
+                autoMigrationResultWithNewAddedColumn.element,
+                autoMigrationResultWithNewAddedColumn
+            )
                 .write(invocation.processingEnv)
 
             invocation.assertCompilationResult {
@@ -131,7 +132,10 @@
                     removedOrRenamedTables = listOf()
                 ),
             )
-            AutoMigrationWriter(mock(XElement::class.java), autoMigrationResultWithNewAddedColumn)
+            AutoMigrationWriter(
+                autoMigrationResultWithNewAddedColumn.element,
+                autoMigrationResultWithNewAddedColumn
+            )
                 .write(invocation.processingEnv)
 
             invocation.assertCompilationResult {