Support line location information in KAPT for top level declarations

This CL fixes a bug in KAPT's MAP_DIAGNOSTIC_LOCATIONS flag where the
location mapping for top level elements didn't work.

I also needed to changes the Kapt3IntegrationTest to copy input files,
otherwise, it cannot map file paths since the mock project has no root
location.

Cleanup tmp folders in KAPT tests, update Kapt3IT test for new lines.

^KT-47934 Fixed
Test: KotlinKapt3IntegrationTests#testErrorLocationMapping

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/Kapt3IT.kt

@@ -479,9 +479,9 @@ open class Kapt3IT : Kapt3BaseIT() {
             val actual = getErrorMessages()
             // try as 0 starting lines first, then as 1 starting line
             try {
-                Assert.assertEquals(genKotlinErrorString(2, 6), actual)
+                Assert.assertEquals(genKotlinErrorString(3, 6), actual)
             } catch (e: AssertionError) {
-                Assert.assertEquals(genKotlinErrorString(3, 7), actual)
+                Assert.assertEquals(genKotlinErrorString(4, 7), actual)
             }
         }
     }

plugins/kapt3/kapt3-base/src/org/jetbrains/kotlin/kapt3/base/stubs/KaptStubLineInformation.kt

@@ -168,6 +168,9 @@ class KaptStubLineInformation {
     }
 
     private fun JCTree.isLocatedInside(declaration: JCTree): Boolean {
+        if (this === declaration) {
+            return true
+        }
         var found = false
 
         declaration.accept(object : TreeScanner() {

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3IntegrationTest.kt

@@ -17,10 +17,15 @@
 package org.jetbrains.kotlin.kapt3.test
 
 import com.intellij.openapi.util.text.StringUtil
+import com.intellij.openapi.vfs.StandardFileSystems
+import com.intellij.psi.PsiManager
 import org.jetbrains.kotlin.base.kapt3.DetectMemoryLeaksMode
 import org.jetbrains.kotlin.base.kapt3.KaptOptions
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSourceLocation
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.codegen.ClassBuilderMode
+import org.jetbrains.kotlin.codegen.CodegenTestFiles
 import org.jetbrains.kotlin.codegen.GenerationUtils
 import org.jetbrains.kotlin.codegen.OriginCollectingClassBuilderFactory
 import org.jetbrains.kotlin.kapt3.AbstractKapt3Extension
@@ -34,13 +39,16 @@ import org.jetbrains.kotlin.kapt3.prettyPrint
 import org.jetbrains.kotlin.kapt3.stubs.ClassFileToSourceStubConverter
 import org.jetbrains.kotlin.kapt3.stubs.ClassFileToSourceStubConverter.KaptStub
 import org.jetbrains.kotlin.kapt3.util.MessageCollectorBackedKaptLogger
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
 import org.jetbrains.kotlin.test.ConfigurationKind
 import org.jetbrains.kotlin.test.KotlinTestUtils
+import org.jetbrains.kotlin.test.util.KtTestUtil
 import org.jetbrains.kotlin.test.util.trimTrailingWhitespacesAndAddNewlineAtEOF
 import org.jetbrains.kotlin.utils.PathUtil
 import java.io.File
 import java.nio.file.Files
+import java.util.ArrayList
 import javax.annotation.processing.Completion
 import javax.annotation.processing.ProcessingEnvironment
 import javax.annotation.processing.Processor
@@ -123,6 +131,24 @@ abstract class AbstractKotlinKapt3IntegrationTest : KotlinKapt3TestBase() {
         }
     }
 
+    override fun loadMultiFiles(files: List<TestFile>) {
+        val project = myEnvironment.project
+        val psiManager = PsiManager.getInstance(project)
+
+        val tmpDir = tmpDir("kaptTest")
+
+        val ktFiles = ArrayList<KtFile>(files.size)
+        for (file in files.sorted()) {
+            if (file.name.endsWith(".kt")) {
+                val tmpKtFile = File(tmpDir, file.name).apply { writeText(file.content) }
+                val virtualFile = StandardFileSystems.local().findFileByPath(tmpKtFile.path) ?: error("Can't find ${file.name}")
+                ktFiles.add(psiManager.findFile(virtualFile) as? KtFile ?: error("Can't load ${file.name}"))
+            }
+        }
+
+        myFiles = CodegenTestFiles.create(ktFiles)
+    }
+
     override fun doMultiFileTest(wholeFile: File, files: List<TestFile>) {
         val txtFile = File(wholeFile.parentFile, wholeFile.nameWithoutExtension + ".it.txt")
 
@@ -167,8 +193,15 @@ abstract class AbstractKotlinKapt3IntegrationTest : KotlinKapt3TestBase() {
         }
     }
 
-    protected inner class Kapt3ExtensionForTests(options: KaptOptions, private val processors: List<Processor>) : AbstractKapt3Extension(
-        options, MessageCollectorBackedKaptLogger(options), compilerConfiguration = myEnvironment.configuration
+    protected inner class Kapt3ExtensionForTests(
+        options: KaptOptions,
+        private val processors: List<Processor>,
+        val messageCollector: LoggingMessageCollector = LoggingMessageCollector()
+    ) : AbstractKapt3Extension(
+        options, MessageCollectorBackedKaptLogger(
+            flags = options,
+            messageCollector = messageCollector
+        ), compilerConfiguration = myEnvironment.configuration
     ) {
         internal var savedStubs: String? = null
         internal var savedBindings: Map<String, KaptJavaFileObject>? = null
@@ -204,4 +237,28 @@ abstract class AbstractKotlinKapt3IntegrationTest : KotlinKapt3TestBase() {
             super.saveIncrementalData(kaptContext, messageCollector, converter)
         }
     }
+
+    class LoggingMessageCollector : MessageCollector {
+        private val _messages = mutableListOf<Message>()
+        val messages: List<Message>
+            get() = _messages
+
+        override fun clear() {
+            _messages.clear()
+        }
+
+        override fun report(severity: CompilerMessageSeverity, message: String, location: CompilerMessageSourceLocation?) {
+            _messages.add(Message(severity, message, location))
+        }
+
+        override fun hasErrors() = _messages.any {
+            it.severity.isError
+        }
+
+        data class Message(
+            val severity: CompilerMessageSeverity,
+            val message: String,
+            val location: CompilerMessageSourceLocation?
+        )
+    }
 }

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/AbstractKotlinKapt3Test.kt

@@ -89,7 +89,7 @@ abstract class AbstractKotlinKapt3Test : KotlinKapt3TestBase() {
         val project = myEnvironment.project
         val psiManager = PsiManager.getInstance(project)
 
-        val tmpDir = KtTestUtil.tmpDir("kaptTest")
+        val tmpDir = tmpDir("kaptTest")
 
         val ktFiles = ArrayList<KtFile>(files.size)
         for (file in files.sorted()) {
@@ -123,7 +123,7 @@ abstract class AbstractKotlinKapt3Test : KotlinKapt3TestBase() {
             projectBaseDir = project.basePath?.let { File(it) }
             compileClasspath.addAll(PathUtil.getJdkClassesRootsFromCurrentJre() + PathUtil.kotlinPathsForIdeaPlugin.stdlibPath)
 
-            sourcesOutputDir = KtTestUtil.tmpDir("kaptRunner")
+            sourcesOutputDir = tmpDir("kaptRunner")
             classesOutputDir = sourcesOutputDir
             stubsOutputDir = sourcesOutputDir
             incrementalDataOutputDir = sourcesOutputDir

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/IrKotlinKaptContextTestGenerated.java

@@ -35,6 +35,11 @@ public class IrKotlinKaptContextTestGenerated extends AbstractIrKotlinKaptContex
         runTest("plugins/kapt3/kapt3-compiler/testData/kotlinRunner/DefaultParameterValues.kt");
     }
 
+    @TestMetadata("ErrorLocationMapping.kt")
+    public void testErrorLocationMapping() throws Exception {
+        runTest("plugins/kapt3/kapt3-compiler/testData/kotlinRunner/ErrorLocationMapping.kt");
+    }
+
     @TestMetadata("NestedClasses.kt")
     public void testNestedClasses() throws Exception {
         runTest("plugins/kapt3/kapt3-compiler/testData/kotlinRunner/NestedClasses.kt");

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/KotlinKapt3IntegrationTests.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.kapt3.test
 
+import junit.framework.TestCase
 import org.jetbrains.kotlin.kapt3.javac.KaptJavaFileObject
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
@@ -29,6 +30,8 @@ import javax.lang.model.element.ElementKind
 import javax.lang.model.element.ExecutableElement
 import javax.lang.model.element.TypeElement
 import javax.lang.model.element.VariableElement
+import javax.lang.model.util.ElementFilter
+import javax.tools.Diagnostic
 import kotlin.system.exitProcess
 
 class KotlinKapt3IntegrationTests : AbstractKotlinKapt3IntegrationTest(), CustomJdkTestLauncher {
@@ -103,6 +106,68 @@ class KotlinKapt3IntegrationTests : AbstractKotlinKapt3IntegrationTest(), Custom
         }
     }
 
+    @Test
+    fun testErrorLocationMapping() {
+        val diagnostics = diagnosticsTest("ErrorLocationMapping", "MyAnnotation") { _, _, processingEnv ->
+            val subject = processingEnv.elementUtils.getTypeElement("Subject")
+            assertNotNull(subject)
+            processingEnv.messager.printMessage(
+                Diagnostic.Kind.NOTE,
+                "note on class",
+                subject
+            )
+            // report error on the field as well
+            val field = ElementFilter.fieldsIn(
+                subject.enclosedElements
+            ).firstOrNull {
+                it.simpleName.toString() == "field"
+            }
+            processingEnv.messager.printMessage(
+                Diagnostic.Kind.NOTE,
+                "note on field",
+                field
+            )
+        }
+        diagnostics.assertContainsDiagnostic("ErrorLocationMapping.kt:1: Note: note on class")
+        diagnostics.assertContainsDiagnostic("ErrorLocationMapping.kt:5: Note: note on field")
+    }
+
+    private fun List<LoggingMessageCollector.Message>.assertContainsDiagnostic(
+        message: String
+    ) {
+        assert(
+            any {
+                it.message.contains(message)
+            }
+        ) {
+            """
+            Didn't find expected diagnostic message.
+            Expected: $message
+            Diagnostics:
+            ${this.joinToString("\n") { "${it.severity}: ${it.message}" }}
+            """.trimIndent()
+        }
+    }
+
+    @Suppress("SameParameterValue")
+    private fun diagnosticsTest(
+        name: String,
+        vararg supportedAnnotations: String,
+        process: (Set<TypeElement>, RoundEnvironment, ProcessingEnvironment) -> Unit
+    ): List<LoggingMessageCollector.Message> {
+        lateinit var messageCollector: LoggingMessageCollector
+        test(
+            name = name,
+            supportedAnnotations = supportedAnnotations
+        ) { typeElements, roundEnv, processingEnv ->
+            val kaptExtension = AnalysisHandlerExtension.getInstances(myEnvironment.project).firstIsInstance<Kapt3ExtensionForTests>()
+            messageCollector = kaptExtension.messageCollector
+            process(typeElements, roundEnv, processingEnv)
+        }
+        return messageCollector.messages
+    }
+
+
     private fun bindingsTest(name: String, test: (File, File, Map<String, KaptJavaFileObject>) -> Unit) {
         test(name, "test.MyAnnotation") { _, _, _ ->
             val kaptExtension = AnalysisHandlerExtension.getInstances(myEnvironment.project).firstIsInstance<Kapt3ExtensionForTests>()

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/KotlinKapt3TestBase.kt

@@ -7,11 +7,13 @@ package org.jetbrains.kotlin.kapt3.test
 
 import org.jetbrains.kotlin.base.kapt3.KaptFlag
 import org.jetbrains.kotlin.codegen.CodegenTestCase
+import org.jetbrains.kotlin.test.util.KtTestUtil
 import java.io.File
 
 abstract class KotlinKapt3TestBase : CodegenTestCase() {
     val kaptFlagsToAdd = mutableListOf<KaptFlag>()
     val kaptFlagsToRemove = mutableListOf<KaptFlag>()
+    private val directoriesToCleanup = mutableListOf<File>()
 
     override fun setUp() {
         super.setUp()
@@ -19,6 +21,11 @@ abstract class KotlinKapt3TestBase : CodegenTestCase() {
         kaptFlagsToRemove.clear()
     }
 
+    override fun tearDown() {
+        directoriesToCleanup.forEach(File::deleteRecursively)
+        super.tearDown()
+    }
+
     protected fun isFlagEnabled(flagName: String, testFile: File): Boolean {
         val stringToCheck = "// $flagName"
         return testFile.useLines { lines -> lines.any { it.trim() == stringToCheck } }
@@ -37,6 +44,10 @@ abstract class KotlinKapt3TestBase : CodegenTestCase() {
         }
     }
 
+    protected fun tmpDir(name: String): File {
+        return KtTestUtil.tmpDir(name).also(directoriesToCleanup::add)
+    }
+
     override fun doTest(filePath: String) {
         val testFile = File(filePath)
 

plugins/kapt3/kapt3-compiler/test/org/jetbrains/kotlin/kapt3/test/KotlinKaptContextTestGenerated.java

@@ -34,6 +34,11 @@ public class KotlinKaptContextTestGenerated extends AbstractKotlinKaptContextTes
         runTest("plugins/kapt3/kapt3-compiler/testData/kotlinRunner/DefaultParameterValues.kt");
     }
 
+    @TestMetadata("ErrorLocationMapping.kt")
+    public void testErrorLocationMapping() throws Exception {
+        runTest("plugins/kapt3/kapt3-compiler/testData/kotlinRunner/ErrorLocationMapping.kt");
+    }
+
     @TestMetadata("NestedClasses.kt")
     public void testNestedClasses() throws Exception {
         runTest("plugins/kapt3/kapt3-compiler/testData/kotlinRunner/NestedClasses.kt");

plugins/kapt3/kapt3-compiler/testData/kotlinRunner/ErrorLocationMapping.it.txt

@@ -0,0 +1,45 @@
+import java.lang.System;
+
+@kotlin.Metadata()
+@java.lang.annotation.Retention(value = java.lang.annotation.RetentionPolicy.RUNTIME)
+public abstract @interface MyAnnotation {
+}
+
+////////////////////
+
+
+import java.lang.System;
+
+@kotlin.Metadata()
+public final class Subject {
+    @org.jetbrains.annotations.NotNull()
+    private final java.lang.String field = "";
+    @org.jetbrains.annotations.NotNull()
+    private final java.lang.String annotationTrigger = "";
+
+    public Subject() {
+        super();
+    }
+
+    @org.jetbrains.annotations.NotNull()
+    public final java.lang.String getField() {
+        return null;
+    }
+
+    @org.jetbrains.annotations.NotNull()
+    public final java.lang.String getAnnotationTrigger() {
+        return null;
+    }
+
+    @MyAnnotation()
+    @java.lang.Deprecated()
+    public static void getAnnotationTrigger$annotations() {
+    }
+}
+
+////////////////////
+
+package error;
+
+public final class NonExistentClass {
+}

plugins/kapt3/kapt3-compiler/testData/kotlinRunner/ErrorLocationMapping.kt

@@ -0,0 +1,11 @@
+class Subject {
+
+
+
+    val field:String = ""
+
+    @MyAnnotation
+    val annotationTrigger: String = ""
+}
+
+internal annotation class MyAnnotation

plugins/kapt3/kapt3-compiler/testData/kotlinRunner/ErrorLocationMapping.txt

@@ -0,0 +1 @@
+