New J2K: Fix converting on copy-past

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

@@ -496,6 +496,7 @@ class JavaToJKTreeBuilder constructor(
     }
 
     private inner class DeclarationMapper(val expressionTreeMapper: ExpressionTreeMapper) {
+
         fun PsiTypeParameterList.toJK(): JKTypeParameterList =
             JKTypeParameterListImpl(typeParameters.map { it.toJK() })
                 .also {
@@ -705,7 +706,8 @@ class JavaToJKTreeBuilder constructor(
 
         fun PsiAnnotation.toJK(): JKAnnotation =
             JKAnnotationImpl(
-                symbolProvider.provideSymbol(nameReferenceElement!!),
+                symbolProvider.provideSymbol(nameReferenceElement!!).safeAs<JKClassSymbol>()
+                    ?: JKUnresolvedClassSymbol(nameReferenceElement!!.text),
                 parameterList.attributes.map { parameter ->
                     if (parameter.nameIdentifier != null) {
                         JKAnnotationNameParameterImpl(
@@ -954,7 +956,14 @@ class JavaToJKTreeBuilder constructor(
     fun buildTree(psi: PsiElement): JKTreeElement? =
         when (psi) {
             is PsiJavaFile -> psi.toJK()
-            else -> error("Cannot convert non-java file")
+            is PsiExpression -> with(expressionTreeMapper) { psi.toJK() }
+            is PsiStatement -> with(declarationMapper) { psi.toJK() }
+            is PsiClass -> with(declarationMapper) { psi.toJK() }
+            is PsiField -> with(declarationMapper) { psi.toJK() }
+            is PsiMethod -> with(declarationMapper) { psi.toJK() }
+            is PsiAnnotation -> with(declarationMapper) { psi.toJK() }
+            else ->
+                null
         }
 
     private val tokenCache = mutableMapOf<PsiElement, JKNonCodeElement>()

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

@@ -18,6 +18,8 @@ package org.jetbrains.kotlin.nj2k
 
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiExpression
+import com.intellij.psi.PsiStatement
 import org.jetbrains.kotlin.j2k.*
 import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
 import org.jetbrains.kotlin.nj2k.tree.prettyDebugPrintTree
@@ -49,7 +51,9 @@ class NewJavaToKotlinConverter(
         symbolProvider.preBuildTree(inputElements)
         val importStorage = ImportStorage()
         val treeBuilder = JavaToJKTreeBuilder(symbolProvider, converterServices, importStorage)
-        val asts = inputElements.mapNotNull(treeBuilder::buildTree)
+        val asts = inputElements.map { element ->
+            element to treeBuilder.buildTree(element)
+        }
 
         val context = ConversionContext(
             symbolProvider,
@@ -57,17 +61,22 @@ class NewJavaToKotlinConverter(
             { it.containingFile in inputElements },
             importStorage
         )
-        ConversionsRunner.doApply(asts, context)
-        val kotlinCodes = asts.map { NewCodeBuilder().run { printCodeOut(it) } }
-        return Result(
-            kotlinCodes.map { code ->
-                ElementResult(
-                    code,
-                    importsToAdd = importStorage.getImports(),
-                    parseContext = ParseContext.TOP_LEVEL
-                )
-            },
-            null
-        )
+
+        ConversionsRunner.doApply(asts.mapNotNull { it.second }, context)
+        val results = asts.map { (element, ast) ->
+            if (ast == null) return@map null
+            val code = NewCodeBuilder().run { printCodeOut(ast) }
+            val parseContext = when (element) {
+                is PsiStatement, is PsiExpression -> ParseContext.CODE_BLOCK
+                else -> ParseContext.TOP_LEVEL
+            }
+            ElementResult(
+                code,
+                importsToAdd = importStorage.getImports(),
+                parseContext = parseContext
+            )
+        }
+
+        return Result(results, null)
     }
 }