New J2K: Add package & imports declarations to JKFile, simplify JKFile constructing

j2k/newSrc/org/jetbrains/kotlin/j2k/JavaToJKTreeBuilder.kt

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.j2k
 
-import com.intellij.lang.java.JavaLanguage
 import com.intellij.psi.*
 import com.intellij.psi.JavaTokenType.SUPER_KEYWORD
 import com.intellij.psi.JavaTokenType.THIS_KEYWORD
@@ -35,6 +34,24 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
 
     private val declarationMapper = DeclarationMapper(expressionTreeMapper)
 
+    private fun PsiJavaFile.toJK(): JKFile =
+        JKFileImpl(
+            packageStatement?.toJK() ?: JKPackageDeclarationImpl(JKNameIdentifierImpl("")),
+            importList?.importStatements?.map { it.toJK() }.orEmpty(),
+            with(declarationMapper) { classes.map { it.toJK() } }
+        )
+
+   private fun PsiPackageStatement.toJK(): JKPackageDeclaration =
+        JKPackageDeclarationImpl(JKNameIdentifierImpl(packageName))
+
+    private fun PsiImportStatement.toJK() =
+        JKImportStatementImpl(
+            JKNameIdentifierImpl(
+                text.substringAfter("import").substringBeforeLast(";").trim()
+            )
+        )
+
+
     private inner class ExpressionTreeMapper {
         fun PsiExpression?.toJK(): JKExpression {
             return when (this) {
@@ -251,9 +268,6 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
             JKTypeParameterImpl(JKNameIdentifierImpl(name!!),
                                 extendsListTypes.map { JKTypeElementImpl(it.toJK(symbolProvider, Nullability.NotNull)) })
 
-        fun PsiPackageStatement.toJK(): JKPackageDeclaration =
-            JKPackageDeclarationImpl(JKNameIdentifierImpl(packageName))
-
         fun PsiClass.toJK(): JKClass {
             val classKind: JKClass.ClassKind = when {
                 isAnnotationType -> JKClass.ClassKind.ANNOTATION
@@ -285,7 +299,13 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
         fun PsiClass.createClassBody() =
             JKClassBodyImpl(
                 children.mapNotNull {
-                    ElementVisitor().apply { it.accept(this) }.resultElement as? JKDeclaration
+                    when (it) {
+                        is PsiEnumConstant -> it.toJK()
+                        is PsiClass -> it.toJK()
+                        is PsiMethod -> it.toJK()
+                        is PsiField -> it.toJK()
+                        else -> null
+                    }
                 }
             )
 
@@ -512,42 +532,11 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
         }.last()
 
 
-       private inner class ElementVisitor : JavaElementVisitor() {
+    fun buildTree(psi: PsiElement): JKTreeElement? =
-
+        when (psi) {
-        var resultElement: JKTreeElement? = null
+            is PsiJavaFile -> psi.toJK()
-
+            else -> error("Cannot convert non-java file")
-        override fun visitClass(aClass: PsiClass) {
-            resultElement = with(declarationMapper) { aClass.toJK() }
         }
 
-        override fun visitPackageStatement(statement: PsiPackageStatement) {
-            resultElement = with(declarationMapper) { statement.toJK() }
-        }
-        override fun visitField(field: PsiField) {
-            resultElement = with(declarationMapper) { field.toJK() }
-        }
-
-        override fun visitEnumConstant(enumConstant: PsiEnumConstant) {
-            resultElement = with(declarationMapper) { enumConstant.toJK() }
-        }
-
-        override fun visitMethod(method: PsiMethod) {
-            resultElement = with(declarationMapper) { method.toJK() }
-        }
-
-        override fun visitFile(file: PsiFile) {
-            resultElement = JKFileImpl().apply {
-                declarationList += file.children.mapNotNull { ElementVisitor().apply { it.accept(this) }.resultElement as? JKDeclaration }
-            }
-        }
-    }
-
-
-    fun buildTree(psi: PsiElement): JKTreeElement? {
-        assert(psi.language.`is`(JavaLanguage.INSTANCE)) { "Unable to build JK Tree using Java Visitor for language ${psi.language}" }
-        val elementVisitor = ElementVisitor()
-        psi.accept(elementVisitor)
-        return elementVisitor.resultElement
-    }
 }
 

j2k/newSrc/org/jetbrains/kotlin/j2k/NewCodeBuilder.kt

@@ -119,27 +119,32 @@ class NewCodeBuilder {
         }
 
         override fun visitFile(file: JKFile) {
-            file.declarationList.find { it is JKPackageDeclaration }?.accept(this)
+            if (file.packageDeclaration.packageName.value.isNotEmpty()) {
-            val collectImportsVisitor = CollectImportsVisitor()
+                file.packageDeclaration.accept(this)
-            file.accept(collectImportsVisitor)
-            collectImportsVisitor.collectedFqNames.forEach {
-                printer.printlnWithNoIndent("import ", it.asString())
-            }
-            file.declarationList.forEach {
-                if (it !is JKPackageDeclaration) it.accept(this)
             }
+            file.importList.forEach { it.accept(this) }
+            file.declarationList.forEach { it.accept(this) }
         }
 
+        private fun String.escapedAsQualifiedName(): String =
+            split('.')
+                .map { it.escaped() }
+                .joinToString(".") { it }
+
         override fun visitPackageDeclaration(packageDeclaration: JKPackageDeclaration) {
             printer.printWithNoIndent("package ")
             val packageNameEscaped =
-                packageDeclaration.packageName.value
+                packageDeclaration.packageName.value.escapedAsQualifiedName()
-                    .split('.')
-                    .map { it.escaped() }
-                    .joinToString(".") { it }
             printer.printlnWithNoIndent(packageNameEscaped)
         }
 
+        override fun visitImportStatement(importStatement: JKImportStatement) {
+            printer.printWithNoIndent("import ")
+            val importNameEscaped =
+                importStatement.name.value.escapedAsQualifiedName()
+            printer.printlnWithNoIndent(importNameEscaped)
+        }
+
         override fun visitBreakStatement(breakStatement: JKBreakStatement) {
             printer.printWithNoIndent("break")
         }
@@ -782,35 +787,6 @@ class NewCodeBuilder {
         ROUND("(", ")"), SQUARE("[", "]"), CURVED("{", "}"), CURVED_MULTILINE("{\n", "}\n"), INLINE_COMMENT("/*", "*/"), ANGLE("<", ">")
     }
 
-    private class CollectImportsVisitor : JKVisitorVoid {
-        val collectedFqNames: MutableList<FqName> = mutableListOf()
-        override fun visitTreeElement(treeElement: JKTreeElement) {
-            treeElement.acceptChildren(this)
-        }
-
-        override fun visitJavaNewExpression(javaNewExpression: JKJavaNewExpression) {
-            collectedFqNames.add(FqName(javaNewExpression.classSymbol.fqName!!))
-            javaNewExpression.acceptChildren(this)
-        }
-
-//        override fun visitTypeElement(typeElement: JKTypeElement, data: Nothing?) {
-//            val classType = typeElement.type as? JKClassType ?: return
-//            collectedFqNames.add(FqName(classType.classReference.fqName!!))
-//        }
-
-        override fun visitClassLiteralExpression(classLiteralExpression: JKClassLiteralExpression) {
-            val type = classLiteralExpression.classType.type
-            if (type is JKClassType) {
-                val fqName = type.classReference.fqName!!
-                val isDefaultImport = fqName.split('.').let {
-                    it.size == 2 && it.first() == "kotlin"
-                }
-                if (!isDefaultImport) {
-                    collectedFqNames.add(FqName(fqName))
-                }
-            }
-        }
-    }
 
     fun printCodeOut(root: JKTreeElement): String {
         Visitor().also { root.accept(it) }

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/impl/jk.kt

@@ -31,10 +31,18 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtNamedFunction
 
-class JKFileImpl : JKFile, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+
+class JKFileImpl(
+    packageDeclaration: JKPackageDeclaration,
+    importList: List<JKImportStatement>,
+    declarationList: List<JKDeclaration>
+) : JKFile, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitFile(this, data)
 
-    override var declarationList by children<JKDeclaration>()
+    override var packageDeclaration: JKPackageDeclaration by child(packageDeclaration)
+    override var importList: List<JKImportStatement> by children(importList)
+    override var declarationList by children(declarationList)
 }
 
 class JKClassImpl(
@@ -506,4 +514,10 @@ class JKClassLiteralExpressionImpl(
     override val classType: JKTypeElement by child(classType)
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitClassLiteralExpression(this, data)
+}
+
+class JKImportStatementImpl(name: JKNameIdentifier) : JKImportStatement, JKBranchElementBase() {
+    override val name: JKNameIdentifier by child(name)
+
+    override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitImportStatement(this, data)
 }

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/jk.kt

@@ -40,7 +40,13 @@ interface PsiOwner {
 
 interface JKDeclaration : JKTreeElement
 
+interface JKImportStatement : JKTreeElement {
+    val name: JKNameIdentifier
+}
+
 interface JKFile : JKTreeElement, JKBranchElement {
+    var packageDeclaration: JKPackageDeclaration
+    var importList: List<JKImportStatement>
     var declarationList: List<JKDeclaration>
 }
 

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/visitors/JKVisitor.kt

@@ -5,6 +5,7 @@ import org.jetbrains.kotlin.j2k.tree.*
 interface JKVisitor<out R, in D> {
     fun visitTreeElement(treeElement: JKTreeElement, data: D): R 
     fun visitDeclaration(declaration: JKDeclaration, data: D): R = visitTreeElement(declaration, data)
+    fun visitImportStatement(importStatement: JKImportStatement, data: D): R = visitTreeElement(importStatement, data)
     fun visitFile(file: JKFile, data: D): R = visitTreeElement(file, data)
     fun visitClass(klass: JKClass, data: D): R = visitDeclaration(klass, data)
     fun visitInheritanceInfo(inheritanceInfo: JKInheritanceInfo, data: D): R = visitTreeElement(inheritanceInfo, data)

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/visitors/JKVisitorVoid.kt

@@ -7,6 +7,8 @@ interface JKVisitorVoid : JKVisitor<Unit, Nothing?> {
     override fun visitTreeElement(treeElement: JKTreeElement, data: Nothing?) = visitTreeElement(treeElement)
     fun visitDeclaration(declaration: JKDeclaration) = visitTreeElement(declaration, null)
     override fun visitDeclaration(declaration: JKDeclaration, data: Nothing?) = visitDeclaration(declaration)
+    fun visitImportStatement(importStatement: JKImportStatement) = visitTreeElement(importStatement, null)
+    override fun visitImportStatement(importStatement: JKImportStatement, data: Nothing?) = visitImportStatement(importStatement)
     fun visitFile(file: JKFile) = visitTreeElement(file, null)
     override fun visitFile(file: JKFile, data: Nothing?) = visitFile(file)
     fun visitClass(klass: JKClass) = visitDeclaration(klass, null)