New J2K: Add imports of used kotlin functions

nj2k/src/org/jetbrains/kotlin/nj2k/ConversionContext.kt

@@ -13,7 +13,8 @@ import org.jetbrains.kotlin.nj2k.tree.JKElement
 data class ConversionContext(
     val symbolProvider: JKSymbolProvider,
     val converter: NewJavaToKotlinConverter,
-    val inConversionContext: (PsiElement) -> Boolean
+    val inConversionContext: (PsiElement) -> Boolean,
+    val importStorage: ImportStorage
 ) {
     val project: Project get() = converter.project
     val typeFlavorCalculator = TypeFlavorCalculator(object : TypeFlavorConverterFacade {
@@ -26,4 +27,5 @@ data class ConversionContext(
 
         override fun inConversionScope(element: PsiElement): Boolean = inConversionContext(element)
     })
+
 }

nj2k/src/org/jetbrains/kotlin/nj2k/ImportStorage.kt

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license 
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.nj2k
+
+import org.jetbrains.kotlin.name.FqName
+
+class ImportStorage {
+    private val imports = mutableSetOf<FqName>()
+
+    fun addImport(import: FqName) {
+        imports += import
+    }
+
+    fun getImports(): Set<FqName> = imports
+}

nj2k/src/org/jetbrains/kotlin/nj2k/JavaToJKTreeBuilder.kt

@@ -39,6 +39,7 @@ import org.jetbrains.kotlin.asJava.elements.KtLightMethod
 import org.jetbrains.kotlin.idea.caches.lightClasses.KtLightClassForDecompiledDeclaration
 import org.jetbrains.kotlin.idea.j2k.IdeaDocCommentConverter
 import org.jetbrains.kotlin.idea.j2k.content
+import org.jetbrains.kotlin.idea.refactoring.fqName.getKotlinFqName
 import org.jetbrains.kotlin.j2k.ReferenceSearcher
 import org.jetbrains.kotlin.j2k.ast.Nullability
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -53,9 +54,10 @@ import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 
-class JavaToJKTreeBuilder(
+class JavaToJKTreeBuilder constructor(
     var symbolProvider: JKSymbolProvider,
-    converterServices: NewJavaToKotlinServices
+    converterServices: NewJavaToKotlinServices,
+    private val importStorage: ImportStorage
 ) {
 
     private val expressionTreeMapper = ExpressionTreeMapper()
@@ -294,12 +296,14 @@ class JavaToJKTreeBuilder(
                         is KtNamedFunction -> {
                             if (origin.isExtensionDeclaration()) {
                                 val receiver = arguments.expressions.firstOrNull()?.toJK()?.parenthesizeIfBinaryExpression()
+                                origin.fqName?.also { importStorage.addImport(it) }
                                 JKJavaMethodCallExpressionImpl(
                                     symbolProvider.provideDirectSymbol(origin) as JKMethodSymbol,
                                     arguments.expressions.drop(1).map { it.toJK() }.toArgumentList(),
                                     typeArguments
                                 ).qualified(receiver)
                             } else {
+                                origin.fqName?.also { importStorage.addImport(it) }
                                 JKJavaMethodCallExpressionImpl(
                                     symbolProvider.provideDirectSymbol(origin) as JKMethodSymbol,
                                     arguments.toJK(),
@@ -308,6 +312,7 @@ class JavaToJKTreeBuilder(
                             }
                         }
                         is KtProperty, is KtPropertyAccessor, is KtParameter -> {
+                            origin.getKotlinFqName()?.also { importStorage.addImport(it) }
                             val property =
                                 if (origin is KtPropertyAccessor) origin.parent as KtProperty
                                 else origin as KtNamedDeclaration

nj2k/src/org/jetbrains/kotlin/nj2k/NewJavaToKotlinConverter.kt

@@ -21,7 +21,6 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.j2k.*
 import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
 import org.jetbrains.kotlin.nj2k.tree.prettyDebugPrintTree
-import org.jetbrains.kotlin.psi.KtPsiFactory
 
 class NewJavaToKotlinConverter(
     val project: Project,
@@ -43,20 +42,25 @@ class NewJavaToKotlinConverter(
     override fun elementsToKotlin(inputElements: List<PsiElement>, processor: WithProgressProcessor): Result {
         val symbolProvider = JKSymbolProvider()
         symbolProvider.preBuildTree(inputElements)
-        val treeBuilder = JavaToJKTreeBuilder(symbolProvider, converterServices)
+        val importStorage = ImportStorage()
+        val treeBuilder = JavaToJKTreeBuilder(symbolProvider, converterServices, importStorage)
         val asts = inputElements.mapNotNull(treeBuilder::buildTree)
-        val factory = KtPsiFactory(project, true)
 
         val context = ConversionContext(
             symbolProvider,
             this,
-            { it.containingFile in inputElements }
+            { it.containingFile in inputElements },
+            importStorage
         )
         ConversionsRunner.doApply(asts, context)
         val kotlinCodes = asts.map { NewCodeBuilder().run { printCodeOut(it) } }
         return Result(
             kotlinCodes.map { code ->
-                ElementResult(code, emptySet(), ParseContext.TOP_LEVEL)
+                ElementResult(
+                    code,
+                    importsToAdd = importStorage.getImports(),
+                    parseContext = ParseContext.TOP_LEVEL
+                )
             },
             null
         )

nj2k/tests/org/jetbrains/kotlin/nj2k/AbstractNewJavaToKotlinConverterSingleFileTest.kt

@@ -20,7 +20,6 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.idea.j2k.IdeaJavaToKotlinServices
 import org.jetbrains.kotlin.j2k.AbstractJavaToKotlinConverterSingleFileTest
 import org.jetbrains.kotlin.j2k.ConverterSettings
-import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.test.KotlinTestUtils
 import java.io.File
 
@@ -34,12 +33,8 @@ abstract class AbstractNewJavaToKotlinConverterSingleFileTest : AbstractJavaToKo
 
     override fun fileToKotlin(text: String, settings: ConverterSettings, project: Project): String {
         val file = createJavaFile(text)
-        val factory = KtPsiFactory(project, true)
-
         return NewJavaToKotlinConverter(project, settings, IdeaJavaToKotlinServices)
-            .filesToKotlin(listOf(file), NewJ2kPostProcessor(true, settings)).let { (results, _) ->
+            .filesToKotlin(listOf(file), NewJ2kPostProcessor(true, settings)).results.single()
-                factory.createFileWithLightClassSupport("Dummy.kt", results.single(), file)
-            }.text
     }
 
     override fun provideExpectedFile(javaPath: String): File =