Kapt3: Generate import statements in stubs

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/KaptContext.kt

@@ -40,6 +40,7 @@ class KaptContext(
     val compiler: KaptJavaCompiler
     val fileManager: JavacFileManager
     val options: Options
+    val javaLog: KaptJavaLog
 
     init {
         KaptJavaLog.preRegister(context, logger.messageCollector)
@@ -50,6 +51,8 @@ class KaptContext(
         fileManager = context.get(JavaFileManager::class.java) as JavacFileManager
         compiler = JavaCompiler.instance(context) as KaptJavaCompiler
 
+        javaLog = compiler.log as KaptJavaLog
+
         options = Options.instance(context)
         for ((key, value) in processorOptions) {
             val option = if (value.isEmpty()) "-A$key" else "-A$key=$value"

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/annotationProcessing.kt

@@ -65,6 +65,7 @@ fun KaptContext.doAnnotationProcessing(
         val parsedJavaFiles = compiler.parseFiles(javaFileObjects)
 
         compilerAfterAP = try {
+            javaLog.interceptorData.files = parsedJavaFiles.map { it.sourceFile to it }.toMap()
             val analyzedFiles = compiler.stopIfErrorOccurred(
                     CompileStates.CompileState.PARSE, compiler.enterTrees(parsedJavaFiles + additionalSources))
             compiler.processAnnotations(analyzedFiles)

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/javac/KaptJavaLog.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.kapt3.javac
 
+import com.sun.tools.javac.tree.JCTree
 import com.sun.tools.javac.util.Context
 import com.sun.tools.javac.util.JCDiagnostic
 import com.sun.tools.javac.util.Log
@@ -24,12 +25,14 @@ import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity.*
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.kapt3.util.MessageCollectorBackedWriter
 import java.io.PrintWriter
+import javax.tools.JavaFileObject
 
 class KaptJavaLog(
         context: Context,
         errWriter: PrintWriter,
         warnWriter: PrintWriter,
-        noticeWriter: PrintWriter
+        noticeWriter: PrintWriter,
+        val interceptorData: DiagnosticInterceptorData
 ) : Log(context, errWriter, warnWriter, noticeWriter) {
     init {
         context.put(Log.outKey, noticeWriter)
@@ -40,9 +43,44 @@ class KaptJavaLog(
             return
         }
 
+        val targetElement = diagnostic.diagnosticPosition
+        if (diagnostic.code.contains("err.cant.resolve") && targetElement != null) {
+            val sourceFile = interceptorData.files[diagnostic.source]
+            if (sourceFile != null) {
+                val insideImports = targetElement.tree in sourceFile.imports
+                // Ignore resolve errors in import statements
+                if (insideImports) return
+            }
+        }
+
         super.report(diagnostic)
     }
 
+    private operator fun <T : JCTree> Iterable<T>.contains(element: JCTree?): Boolean {
+        if (element == null) {
+            return false
+        }
+
+        var found = false
+        val visitor = object : JCTree.Visitor() {
+            override fun visitImport(that: JCTree.JCImport) {
+                super.visitImport(that)
+                if (!found) that.qualid.accept(this)
+            }
+
+            override fun visitSelect(that: JCTree.JCFieldAccess) {
+                super.visitSelect(that)
+                if (!found) that.selected.accept(this)
+            }
+
+            override fun visitTree(that: JCTree) {
+                if (!found && element == that) found = true
+            }
+        }
+        this.forEach { if (!found) it.accept(visitor) }
+        return found
+    }
+
     companion object {
         private val IGNORED_DIAGNOSTICS = setOf(
                 "compiler.err.name.clash.same.erasure",
@@ -50,16 +88,22 @@ class KaptJavaLog(
                 "compiler.err.name.clash.same.erasure.no.override.1",
                 "compiler.err.name.clash.same.erasure.no.hide",
                 "compiler.err.already.defined",
-                "compiler.err.annotation.type.not.applicable")
+                "compiler.err.annotation.type.not.applicable",
+                "compiler.err.doesnt.exist")
 
         internal fun preRegister(context: Context, messageCollector: MessageCollector) {
+            val interceptorData = DiagnosticInterceptorData()
             context.put(Log.logKey, Context.Factory<Log> {
                 fun makeWriter(severity: CompilerMessageSeverity) = PrintWriter(MessageCollectorBackedWriter(messageCollector, severity))
                 val errWriter = makeWriter(ERROR)
                 val warnWriter = makeWriter(STRONG_WARNING)
                 val noticeWriter = makeWriter(INFO)
-                KaptJavaLog(it, errWriter, warnWriter, noticeWriter)
+                KaptJavaLog(it, errWriter, warnWriter, noticeWriter, interceptorData)
             })
         }
     }
+
+    class DiagnosticInterceptorData {
+        var files: Map<JavaFileObject, JCTree.JCCompilationUnit> = emptyMap()
+    }
 }

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/javac/KaptTreeMaker.kt

@@ -26,7 +26,7 @@ import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.Type.*
 
 class KaptTreeMaker(context: Context) : TreeMaker(context) {
-    private val nameTable = Names.instance(context).table
+    val nameTable = Names.instance(context).table
 
     fun Type(type: Type): JCTree.JCExpression {
         convertBuiltinType(type)?.let { return it }

plugins/kapt3/src/org/jetbrains/kotlin/kapt3/stubs/ClassFileToSourceStubConverter.kt

@@ -29,7 +29,9 @@ import org.jetbrains.kotlin.kapt3.*
 import org.jetbrains.kotlin.kapt3.javac.KaptTreeMaker
 import org.jetbrains.kotlin.kapt3.javac.KaptJavaFileObject
 import org.jetbrains.kotlin.kapt3.util.*
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassNotAny
 import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassOrAny
@@ -131,7 +133,7 @@ class ClassFileToSourceStubConverter(
 
         val classDeclaration = convertClass(clazz, packageName, true) ?: return null
 
-        val imports = JavacList.nil<JCTree>()
+        val imports = if (correctErrorTypes) convertImports(ktFile, classDeclaration) else JavacList.nil()
         val classes = JavacList.of<JCTree>(classDeclaration)
 
         val topLevel = treeMaker.TopLevel(packageAnnotations, packageClause, imports + classes)
@@ -144,6 +146,44 @@ class ClassFileToSourceStubConverter(
         return topLevel
     }
 
+    private fun convertImports(file: KtFile, classDeclaration: JCClassDecl): JavacList<JCTree> {
+        val imports = mutableListOf<JCTree>()
+
+        for (importDirective in file.importDirectives) {
+            // Qualified name should be valid Java fq-name
+            val importedFqName = importDirective.importedFqName ?: continue
+            if (!isValidQualifiedName(importedFqName)) continue
+
+            val shortName = importedFqName.shortName()
+            if (shortName.asString() == classDeclaration.simpleName.toString()) continue
+
+            // If alias is specified, it also should be valid Java name
+            val aliasName = importDirective.aliasName
+            if (aliasName != null /*TODO support aliases */ /*&& getValidIdentifierName(aliasName) == null*/) continue
+
+            val importedReference = getReferenceExpression(importDirective.importedReference)
+                    ?.let { kaptContext.bindingContext[BindingContext.REFERENCE_TARGET, it] }
+
+            if (importedReference is CallableDescriptor) continue
+
+            val importedExpr = treeMaker.FqName(importedFqName.asString())
+
+            imports += if (importDirective.isAllUnder) {
+                treeMaker.Import(treeMaker.Select(importedExpr, treeMaker.nameTable.names.asterisk), false)
+            } else {
+                treeMaker.Import(importedExpr, false)
+            }
+        }
+
+        return JavacList.from(imports)
+    }
+
+    private tailrec fun getReferenceExpression(expression: KtExpression?): KtReferenceExpression? = when (expression) {
+        is KtReferenceExpression -> expression
+        is KtQualifiedExpression -> getReferenceExpression(expression.selectorExpression)
+        else -> null
+    }
+
     /**
      * Returns false for the inner classes or if the origin for the class was not found.
      */
@@ -456,6 +496,10 @@ class ClassFileToSourceStubConverter(
         return ifNonError()
     }
 
+    private fun isValidQualifiedName(name: FqName): Boolean {
+        return name.pathSegments().all { getValidIdentifierName(it.asString(), false) != null }
+    }
+
     fun getValidIdentifierName(name: String, canBeConstructor: Boolean = false): String? {
         if (canBeConstructor && name == "<init>") {
             return name

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

@@ -96,6 +96,7 @@ abstract class AbstractClassFileToSourceStubConverterTest : AbstractKotlinKapt3T
 
         val javaFiles = convert(kaptRunner, typeMapper, generateNonExistentClass, correctErrorTypes)
 
+        kaptRunner.javaLog.interceptorData.files = javaFiles.map { it.sourceFile to it }.toMap()
         if (validate) kaptRunner.compiler.enterTrees(javaFiles)
 
         val actualRaw = javaFiles.sortedBy { it.sourceFile.name }.joinToString (FILE_SEPARATOR)

plugins/kapt3/test/org/jetbrains/kotlin/kapt3/test/ClassFileToSourceStubConverterTestGenerated.java

@@ -102,6 +102,12 @@ public class ClassFileToSourceStubConverterTestGenerated extends AbstractClassFi
         doTest(fileName);
     }
 
+    @TestMetadata("importsForErrorTypes.kt")
+    public void testImportsForErrorTypes() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("plugins/kapt3/testData/converter/importsForErrorTypes.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("inheritanceSimple.kt")
     public void testInheritanceSimple() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("plugins/kapt3/testData/converter/inheritanceSimple.kt");

plugins/kapt3/testData/converter/importsForErrorTypes.kt

@@ -0,0 +1,12 @@
+// CORRECT_ERROR_TYPES
+
+import java.util.zip.*
+import java.util.Date
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.TimeUnit.MICROSECONDS
+import java.util.concurrent.TimeUnit.*
+import java.util.Calendar.*
+
+fun test(): Any? {
+    return null
+}

plugins/kapt3/testData/converter/importsForErrorTypes.txt

@@ -0,0 +1,18 @@
+import java.util.zip.*;
+import java.util.Date;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeUnit.MICROSECONDS;
+import java.util.concurrent.TimeUnit.*;
+import java.util.Calendar.*;
+
+public final class ImportsForErrorTypesKt {
+
+    public ImportsForErrorTypesKt() {
+        super();
+    }
+
+    @org.jetbrains.annotations.Nullable()
+    public static final java.lang.Object test() {
+        return null;
+    }
+}

plugins/kapt3/testData/converter/nonExistentClassTypesConversion.txt

@@ -1,3 +1,5 @@
+import java.util.Calendar;
+
 public final class MyType<T extends java.lang.Object> {
 
     public MyType() {
@@ -8,6 +10,8 @@ public final class MyType<T extends java.lang.Object> {
 ////////////////////
 
 
+import java.util.Calendar;
+
 @kotlin.Suppress(names = {"UNRESOLVED_REFERENCE"})
 public final class Test {
     @org.jetbrains.annotations.NotNull()