From f28f340840f8b477bacdac8f02f36967b967d369 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Fri, 29 Mar 2019 15:08:57 +0300
Subject: [PATCH] New J2K: Wrap ast to a JKTreeRoot

---
 nj2k/src/org/jetbrains/kotlin/nj2k/ConversionsRunner.kt   | 2 +-
 nj2k/src/org/jetbrains/kotlin/nj2k/NewCodeBuilder.kt      | 4 ++++
 nj2k/src/org/jetbrains/kotlin/nj2k/tree/impl/jk.kt        | 4 ++++
 nj2k/src/org/jetbrains/kotlin/nj2k/tree/jk.kt             | 4 ++++
 .../org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitor.kt  | 1 +
 .../jetbrains/kotlin/nj2k/tree/visitors/JKVisitorVoid.kt  | 2 ++
 .../nj2k/tree/visitors/JKVisitorWithCommentsPrinting.kt   | 8 ++++++++
 7 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/nj2k/src/org/jetbrains/kotlin/nj2k/ConversionsRunner.kt b/nj2k/src/org/jetbrains/kotlin/nj2k/ConversionsRunner.kt
index 963bf8ae579..d381fcb255e 100644
--- a/nj2k/src/org/jetbrains/kotlin/nj2k/ConversionsRunner.kt
+++ b/nj2k/src/org/jetbrains/kotlin/nj2k/ConversionsRunner.kt
@@ -85,7 +85,7 @@ object ConversionsRunner {
             +SortClassMembersConversion()
         }
 
-    fun doApply(trees: List<JKTreeElement>, context: ConversionContext) {
+    fun doApply(trees: List<JKTreeRoot>, context: ConversionContext) {
         val conversion = createRootConversion(context)
         conversion.runConversion(trees, context)
     }
diff --git a/nj2k/src/org/jetbrains/kotlin/nj2k/NewCodeBuilder.kt b/nj2k/src/org/jetbrains/kotlin/nj2k/NewCodeBuilder.kt
index 6bc95876b7e..9d79309f0a6 100644
--- a/nj2k/src/org/jetbrains/kotlin/nj2k/NewCodeBuilder.kt
+++ b/nj2k/src/org/jetbrains/kotlin/nj2k/NewCodeBuilder.kt
@@ -116,6 +116,10 @@ class NewCodeBuilder {
             printer.print("/* !!! Hit visitElement for element type: ${treeElement::class} !!! */")
         }
 
+        override fun visitTreeRootRaw(treeRoot: JKTreeRoot) {
+            treeRoot.element.accept(this)
+        }
+
         override fun visitKtTryExpressionRaw(ktTryExpression: JKKtTryExpression) {
             printer.printWithNoIndent("try ")
             ktTryExpression.tryBlock.accept(this)
diff --git a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/impl/jk.kt b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/impl/jk.kt
index 652944c503d..640a814840e 100644
--- a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/impl/jk.kt
+++ b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/impl/jk.kt
@@ -29,6 +29,10 @@ import org.jetbrains.kotlin.nj2k.tree.JKLiteralExpression.LiteralType.NULL
 import org.jetbrains.kotlin.nj2k.tree.visitors.JKVisitor
 import org.jetbrains.kotlin.psi.KtClass
 
+class JKTreeRootImpl(element: JKTreeElement) : JKTreeRoot, JKBranchElementBase() {
+    override var element by child(element)
+    override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitTreeRoot(this, data)
+}
 
 class JKFileImpl(
     packageDeclaration: JKPackageDeclaration,
diff --git a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/jk.kt b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/jk.kt
index d505cbb0769..4ca9da20301 100644
--- a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/jk.kt
+++ b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/jk.kt
@@ -32,6 +32,10 @@ interface JKTreeElement : JKElement, JKNonCodeElementsListOwner {
     fun copy(): JKTreeElement
 }
 
+interface JKTreeRoot : JKTreeElement {
+    var element: JKTreeElement
+}
+
 interface PsiOwner {
     var psi: PsiElement?
 }
diff --git a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitor.kt b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitor.kt
index e8516684909..2f4a70566f8 100644
--- a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitor.kt
+++ b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitor.kt
@@ -4,6 +4,7 @@ import org.jetbrains.kotlin.nj2k.tree.*
 
 interface JKVisitor<out R, in D> {
     fun visitTreeElement(treeElement: JKTreeElement, data: D): R 
+    fun visitTreeRoot(treeRoot: JKTreeRoot, data: D): R = visitTreeElement(treeRoot, data)
     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)
diff --git a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitorVoid.kt b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitorVoid.kt
index a67de1c0af0..99abbb802c4 100644
--- a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitorVoid.kt
+++ b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitorVoid.kt
@@ -5,6 +5,8 @@ import org.jetbrains.kotlin.nj2k.tree.*
 interface JKVisitorVoid : JKVisitor<Unit, Nothing?> {
     fun visitTreeElement(treeElement: JKTreeElement) 
     override fun visitTreeElement(treeElement: JKTreeElement, data: Nothing?) = visitTreeElement(treeElement)
+    fun visitTreeRoot(treeRoot: JKTreeRoot) = visitTreeElement(treeRoot, null)
+    override fun visitTreeRoot(treeRoot: JKTreeRoot, data: Nothing?) = visitTreeRoot(treeRoot)
     fun visitDeclaration(declaration: JKDeclaration) = visitTreeElement(declaration, null)
     override fun visitDeclaration(declaration: JKDeclaration, data: Nothing?) = visitDeclaration(declaration)
     fun visitImportStatement(importStatement: JKImportStatement) = visitTreeElement(importStatement, null)
diff --git a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitorWithCommentsPrinting.kt b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitorWithCommentsPrinting.kt
index 983784faaf2..608a509e6f9 100644
--- a/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitorWithCommentsPrinting.kt
+++ b/nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitorWithCommentsPrinting.kt
@@ -14,6 +14,14 @@ interface JKVisitorWithCommentsPrinting : JKVisitorVoid {
 
     fun visitTreeElementRaw(treeElement: JKTreeElement) 
 
+    override fun visitTreeRoot(treeRoot: JKTreeRoot) {
+        printLeftNonCodeElements(treeRoot)
+        visitTreeRootRaw(treeRoot)
+        printRightNonCodeElements(treeRoot)
+    }
+
+    fun visitTreeRootRaw(treeRoot: JKTreeRoot) = visitTreeElementRaw(treeRoot)
+
     override fun visitDeclaration(declaration: JKDeclaration) {
         printLeftNonCodeElements(declaration)
         visitDeclarationRaw(declaration)