New J2K: Add AST copy tree function

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

@@ -5,16 +5,12 @@
 
 package org.jetbrains.kotlin.j2k
 
-import org.jetbrains.kotlin.j2k.tree.JKArrayAccessExpression
 import org.jetbrains.kotlin.j2k.tree.JKElement
-import org.jetbrains.kotlin.j2k.tree.JKFieldAccessExpression
+import org.jetbrains.kotlin.j2k.tree.impl.JKElementBase
-import org.jetbrains.kotlin.j2k.tree.impl.JKArrayAccessExpressionImpl
-import org.jetbrains.kotlin.j2k.tree.impl.JKFieldAccessExpressionImpl
 
-fun <T : JKElement> T.copyTree(): T {
+fun <T : JKElement> T.copyTree(): T =
-    return when (this) {
+    when (this) {
-        is JKFieldAccessExpression -> JKFieldAccessExpressionImpl(identifier)
+        is JKElementBase ->
-        is JKArrayAccessExpression -> JKArrayAccessExpressionImpl(expression.copyTree(), indexExpression.copyTree())
+            this.copy() as T
-        else -> TODO("Not supported ${this::class}")
+        else -> TODO("Not supported+$this.toString()")
-    } as T
+    }
-}

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

@@ -37,7 +37,7 @@ private class JKListChild<T : JKElement>(val value: Int) : ReadWriteProperty<JKB
     }
 }
 
-abstract class JKElementBase : JKTreeElement {
+abstract class JKElementBase : JKTreeElement, Cloneable {
     override var parent: JKElement? = null
 
     final override fun detach(from: JKElement) {
@@ -64,6 +64,9 @@ abstract class JKElementBase : JKTreeElement {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitTreeElement(this, data)
 
     override fun <D> acceptChildren(visitor: JKVisitor<Unit, D>, data: D) {}
+
+    override fun copy(): JKTreeElement =
+        clone() as JKTreeElement
 }
 
 abstract class JKBranchElementBase : JKElementBase(), JKBranchElement {
@@ -110,5 +113,35 @@ abstract class JKBranchElementBase : JKElementBase(), JKBranchElement {
         check(valid)
     }
 
-    override val children: MutableList<Any> = mutableListOf()
+    final override var children: MutableList<Any> = mutableListOf()
+        private set
+
+    override fun copy(): JKTreeElement {
+        val cloned = super.copy() as JKBranchElementBase
+        val deepClonedChildren =
+            cloned.children.map {
+                when (it) {
+                    is JKElementBase -> it.copy()
+                    is List<*> -> (it as List<JKTreeElement>).map { it.copy() }
+                    else -> error("Tree is corrupted")
+                }
+            }
+
+        deepClonedChildren.forEach { child ->
+            when (child) {
+                is JKElementBase -> {
+                    child.detach(this)
+                    child.attach(cloned)
+                }
+                is List<*> -> (child as List<JKTreeElement>).forEach {
+                    it.detach(this)
+                    it.attach(cloned)
+                }
+            }
+        }
+        cloned.children = deepClonedChildren.toMutableList()
+        return cloned
+    }
+
+
 }

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

@@ -228,6 +228,8 @@ class JKStubExpressionImpl : JKStubExpression, JKElementBase() {
 }
 
 object JKBodyStub : JKBlock, JKTreeElement {
+    override fun copy(): JKTreeElement = this
+
     override var statements: List<JKStatement>
         get() = emptyList()
         set(value) {}

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

@@ -30,6 +30,8 @@ interface JKTreeElement : JKElement {
     fun <D> acceptChildren(visitor: JKVisitor<Unit, D>, data: D)
 
     fun acceptChildren(visitor: JKVisitor<Unit, Nothing?>) = acceptChildren(visitor, null)
+
+    fun copy(): JKTreeElement
 }
 
 interface JKDeclaration : JKTreeElement