Uast: updates for consistency tests and fixes for missed things

we should use `sourcePsi` instead of `psi` for parents matching, because light `psi` will not be tested, and it is not enough to store only PSI as map key to match parent consistency, because some different UElements could have same sourcePsi.

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt

@@ -153,6 +153,12 @@ fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
 
     val result = KotlinUastLanguagePlugin().convertElementWithParent(parentUnwrapped, null)
 
+    if (result is KotlinUBlockExpression && element is UClass) {
+        return KotlinUDeclarationsExpression(result).apply {
+            declarations = listOf(element)
+        }
+    }
+
     if (result is UEnumConstant && element is UDeclaration) {
         return result.initializingClass
     }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt

@@ -550,7 +550,9 @@ private fun convertVariablesDeclaration(
         psi: KtVariableDeclaration,
         parent: UElement?
 ): UDeclarationsExpression {
-    val declarationsExpression = KotlinUDeclarationsExpression(null, parent, psi)
+    val declarationsExpression = parent as? KotlinUDeclarationsExpression
+            ?: psi.parent.toUElementOfType<UDeclarationsExpression>() as? KotlinUDeclarationsExpression
+            ?: KotlinUDeclarationsExpression(null, parent, psi)
     val parentPsiElement = parent?.psi
     val variable = KotlinUAnnotatedLocalVariable(
             UastKotlinPsiVariable.create(psi, parentPsiElement, declarationsExpression), psi, declarationsExpression) { annotationParent ->

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt.173

@@ -545,7 +545,9 @@ private fun convertVariablesDeclaration(
         psi: KtVariableDeclaration,
         parent: UElement?
 ): UDeclarationsExpression {
-    val declarationsExpression = KotlinUDeclarationsExpression(null, parent, psi)
+    val declarationsExpression = parent as? KotlinUDeclarationsExpression
+            ?: psi.parent.toUElementOfType<UDeclarationsExpression>() as? KotlinUDeclarationsExpression
+            ?: KotlinUDeclarationsExpression(null, parent, psi)
     val parentPsiElement = parent?.psi
     val variable = KotlinUAnnotatedLocalVariable(
             UastKotlinPsiVariable.create(psi, parentPsiElement, declarationsExpression), psi, declarationsExpression) { annotationParent ->

plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt

@@ -13,6 +13,7 @@ import org.jetbrains.uast.UFile
 import org.jetbrains.uast.kotlin.JvmDeclarationUElementPlaceholder
 import org.jetbrains.uast.kotlin.KOTLIN_CACHED_UELEMENT_KEY
 import org.jetbrains.uast.kotlin.KotlinUastLanguagePlugin
+import org.jetbrains.uast.sourcePsiElement
 import org.jetbrains.uast.test.common.RenderLogTestBase
 import org.jetbrains.uast.visitor.UastVisitor
 import org.junit.Assert
@@ -39,7 +40,14 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog
     }
 
     private fun checkParentConsistency(file: UFile) {
-        val parentMap = mutableMapOf<PsiElement, String>()
+        val parentMap = mutableMapOf<PsiElement, MutableMap<String, String>>()
+
+        operator fun MutableMap<PsiElement, MutableMap<String, String>>.get(psi: PsiElement, cls: String?) =
+            parentMap.getOrPut(psi) { mutableMapOf() }[cls]
+
+        operator fun MutableMap<PsiElement, MutableMap<String, String>>.set(psi: PsiElement, cls: String, v: String) {
+            parentMap.getOrPut(psi) { mutableMapOf() }[cls] = v
+        }
 
         file.accept(object : UastVisitor {
             private val parentStack = Stack<UElement>()
@@ -52,10 +60,8 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog
                 else {
                     Assert.assertEquals("Wrong parent of $node", parentStack.peek(), parent)
                 }
-                node.psi?.let {
+                node.sourcePsiElement?.let {
-                    if (it !in parentMap) {
+                    parentMap[it, node.asLogString()] = parentStack.reversed().joinToString { it.asLogString() }
-                        parentMap[it] = parentStack.reversed().joinToString { it.asLogString() }
-                    }
                 }
                 parentStack.push(node)
                 return false
@@ -72,7 +78,7 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog
         file.psi.accept(object : PsiRecursiveElementVisitor() {
             override fun visitElement(element: PsiElement) {
                 val uElement = KotlinUastLanguagePlugin().convertElementWithParent(element, null)
-                val expectedParents = parentMap[element]
+                val expectedParents = parentMap[element, uElement?.asLogString()]
                 if (expectedParents != null) {
                     assertNotNull("Expected to be able to convert PSI element $element", uElement)
                     val parents = generateSequence(uElement!!.uastParent) { it.uastParent }.joinToString { it.asLogString() }

plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt.173

@@ -13,6 +13,7 @@ import org.jetbrains.uast.UFile
 import org.jetbrains.uast.kotlin.JvmDeclarationUElementPlaceholder
 import org.jetbrains.uast.kotlin.KOTLIN_CACHED_UELEMENT_KEY
 import org.jetbrains.uast.kotlin.KotlinUastLanguagePlugin
+import org.jetbrains.uast.sourcePsiElement
 import org.jetbrains.uast.test.common.RenderLogTestBase
 import org.jetbrains.uast.visitor.UastVisitor
 import org.junit.Assert
@@ -39,7 +40,14 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog
     }
 
     private fun checkParentConsistency(file: UFile) {
-        val parentMap = mutableMapOf<PsiElement, String>()
+        val parentMap = mutableMapOf<PsiElement, MutableMap<String, String>>()
+
+        operator fun MutableMap<PsiElement, MutableMap<String, String>>.get(psi: PsiElement, cls: String?) =
+            parentMap.getOrPut(psi) { mutableMapOf() }[cls]
+
+        operator fun MutableMap<PsiElement, MutableMap<String, String>>.set(psi: PsiElement, cls: String, v: String) {
+            parentMap.getOrPut(psi) { mutableMapOf() }[cls] = v
+        }
 
         file.accept(object : UastVisitor {
             private val parentStack = Stack<UElement>()
@@ -52,10 +60,8 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog
                 else {
                     Assert.assertEquals("Wrong parent of $node", parentStack.peek(), parent)
                 }
-                node.psi?.let {
+                node.sourcePsiElement?.let {
-                    if (it !in parentMap) {
+                    parentMap[it, node.asLogString()] = parentStack.reversed().joinToString { it.asLogString() }
-                        parentMap[it] = parentStack.reversed().joinToString { it.asLogString() }
-                    }
                 }
                 parentStack.push(node)
                 return false
@@ -72,7 +78,7 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog
         file.psi.accept(object : PsiRecursiveElementVisitor() {
             override fun visitElement(element: PsiElement) {
                 val uElement = KotlinUastLanguagePlugin().convertElementWithParent(element, null)
-                val expectedParents = parentMap[element]
+                val expectedParents = parentMap[element, uElement?.asLogString()]
                 if (expectedParents != null) {
                     assertNotNull("Expected to be able to convert PSI element $element", uElement)
                     val parents = generateSequence(uElement!!.uastParent) { it.uastParent }.joinToString { it.asLogString() }