From dc73d415ad9a1ebbd6d071236ca3f25ba9114bfe Mon Sep 17 00:00:00 2001
From: Simon Ogorodnik <Simon.Ogorodnik@jetbrains.com>
Date: Mon, 20 Nov 2017 23:11:06 +0300
Subject: [PATCH] New J2K: Fix declarations to JK conversion

---
 .../org/jetbrains/kotlin/j2k/JavaToJKTreeBuilder.kt   | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/j2k/newSrc/org/jetbrains/kotlin/j2k/JavaToJKTreeBuilder.kt b/j2k/newSrc/org/jetbrains/kotlin/j2k/JavaToJKTreeBuilder.kt
index 7dc081c2066..b7ef75dce08 100644
--- a/j2k/newSrc/org/jetbrains/kotlin/j2k/JavaToJKTreeBuilder.kt
+++ b/j2k/newSrc/org/jetbrains/kotlin/j2k/JavaToJKTreeBuilder.kt
@@ -186,7 +186,9 @@ class JavaToJKTreeBuilder {
                 isInterface -> JKClass.ClassKind.INTERFACE
                 else -> JKClass.ClassKind.CLASS
             }
-            return JKClassImpl(with(ModifierMapper) { modifierList.toJK() }, JKNameIdentifierImpl(this.name!!), classKind)
+            return JKClassImpl(with(ModifierMapper) { modifierList.toJK() }, JKNameIdentifierImpl(this.name!!), classKind).apply {
+                declarations = children.mapNotNull { ElementVisitor().apply { it.accept(this) }.resultElement as? JKDeclaration }
+            }
         }
 
 
@@ -266,10 +268,6 @@ class JavaToJKTreeBuilder {
 
     private class ElementVisitor : JavaElementVisitor() {
 
-        override fun visitElement(element: PsiElement) {
-            element.acceptChildren(this)
-        }
-
         var resultElement: JKElement? = null
 
         override fun visitClass(aClass: PsiClass) {
@@ -280,6 +278,9 @@ class JavaToJKTreeBuilder {
             resultElement = with(DeclarationMapper) { field.toJK() }
         }
 
+        override fun visitFile(file: PsiFile) {
+            file.acceptChildren(this)
+        }
     }