diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt index 3b2d33310e4..204b32fed2e 100644 --- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt +++ b/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 } diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt index 706ed14eaf6..ec847267a70 100644 --- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt +++ b/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() as? KotlinUDeclarationsExpression + ?: KotlinUDeclarationsExpression(null, parent, psi) val parentPsiElement = parent?.psi val variable = KotlinUAnnotatedLocalVariable( UastKotlinPsiVariable.create(psi, parentPsiElement, declarationsExpression), psi, declarationsExpression) { annotationParent -> diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt.173 b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt.173 index 6ab7aed3629..19e19208ff7 100644 --- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt.173 +++ b/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() as? KotlinUDeclarationsExpression + ?: KotlinUDeclarationsExpression(null, parent, psi) val parentPsiElement = parent?.psi val variable = KotlinUAnnotatedLocalVariable( UastKotlinPsiVariable.create(psi, parentPsiElement, declarationsExpression), psi, declarationsExpression) { annotationParent -> diff --git a/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt b/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt index 7d71ce2970a..f65e2b0aac7 100644 --- a/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt +++ b/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() + val parentMap = mutableMapOf>() + + operator fun MutableMap>.get(psi: PsiElement, cls: String?) = + parentMap.getOrPut(psi) { mutableMapOf() }[cls] + + operator fun MutableMap>.set(psi: PsiElement, cls: String, v: String) { + parentMap.getOrPut(psi) { mutableMapOf() }[cls] = v + } file.accept(object : UastVisitor { private val parentStack = Stack() @@ -52,10 +60,8 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog else { Assert.assertEquals("Wrong parent of $node", parentStack.peek(), parent) } - node.psi?.let { - if (it !in parentMap) { - parentMap[it] = parentStack.reversed().joinToString { it.asLogString() } - } + node.sourcePsiElement?.let { + parentMap[it, node.asLogString()] = 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() } diff --git a/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt.173 b/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt.173 index f706d21b428..23f02ba1b5b 100644 --- a/plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt.173 +++ b/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() + val parentMap = mutableMapOf>() + + operator fun MutableMap>.get(psi: PsiElement, cls: String?) = + parentMap.getOrPut(psi) { mutableMapOf() }[cls] + + operator fun MutableMap>.set(psi: PsiElement, cls: String, v: String) { + parentMap.getOrPut(psi) { mutableMapOf() }[cls] = v + } file.accept(object : UastVisitor { private val parentStack = Stack() @@ -52,10 +60,8 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog else { Assert.assertEquals("Wrong parent of $node", parentStack.peek(), parent) } - node.psi?.let { - if (it !in parentMap) { - parentMap[it] = parentStack.reversed().joinToString { it.asLogString() } - } + node.sourcePsiElement?.let { + parentMap[it, node.asLogString()] = 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() }