UastKotlinPsiVariable psiParent made lazy to eliminate recursion when building UAST parents

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/psi/UastKotlinPsiVariable.kt

@@ -11,12 +11,8 @@ import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
-import org.jetbrains.uast.*
 import org.jetbrains.kotlin.types.isError
-import org.jetbrains.uast.UDeclaration
+import org.jetbrains.uast.*
-import org.jetbrains.uast.UElement
-import org.jetbrains.uast.UastErrorType
-import org.jetbrains.uast.getParentOfType
 import org.jetbrains.uast.kotlin.analyze
 import org.jetbrains.uast.kotlin.lz
 import org.jetbrains.uast.kotlin.orAnonymous
@@ -27,7 +23,7 @@ class UastKotlinPsiVariable private constructor(
         name: String,
         typeProducer: () -> PsiType,
         val ktInitializer: KtExpression?,
-        val psiParent: PsiElement?,
+        psiParentProducer: () -> PsiElement?,
         val containingElement: UElement,
         val ktElement: KtElement
 ) : LightVariableBuilder(
@@ -37,6 +33,8 @@ class UastKotlinPsiVariable private constructor(
         KotlinLanguage.INSTANCE
 ), PsiLocalVariable {
 
+    val psiParent by lz(psiParentProducer)
+
     private val psiType: PsiType by lz(typeProducer)
 
     private val psiTypeElement: PsiTypeElement by lz {
@@ -86,46 +84,37 @@ class UastKotlinPsiVariable private constructor(
                     name = declaration.name.orAnonymous("<unnamed>"),
                     typeProducer = { declaration.getType(containingElement) ?: UastErrorType },
                     ktInitializer = initializerExpression,
-                    psiParent = psiParent,
+                    psiParentProducer = { psiParent },
-                    containingElement = containingElement,
-                    ktElement = declaration)
-        }
-        
-        fun create(declaration: KtDestructuringDeclaration, containingElement: UElement): PsiLocalVariable {
-            val psiParent = containingElement.getParentOfType<UDeclaration>()?.psi ?: declaration.parent
-            return UastKotlinPsiVariable(
-                    manager = declaration.manager,
-                    name = "var" + Integer.toHexString(declaration.getHashCode()),
-                    typeProducer = { declaration.getType(containingElement) ?: UastErrorType },
-                    ktInitializer = declaration.initializer,
-                    psiParent = psiParent,
                     containingElement = containingElement,
                     ktElement = declaration)
         }
 
-        fun create(initializer: KtExpression, containingElement: UElement, parent: PsiElement): PsiLocalVariable {
+        fun create(declaration: KtDestructuringDeclaration, containingElement: UElement): PsiLocalVariable = UastKotlinPsiVariable(
-            val psiParent = containingElement.getParentOfType<UDeclaration>()?.psi ?: parent
+                manager = declaration.manager,
-            return UastKotlinPsiVariable(
+                name = "var" + Integer.toHexString(declaration.getHashCode()),
-                    manager = initializer.manager,
+                typeProducer = { declaration.getType(containingElement) ?: UastErrorType },
-                    name = "var" + Integer.toHexString(initializer.getHashCode()),
+                ktInitializer = declaration.initializer,
-                    typeProducer = { initializer.getType(containingElement) ?: UastErrorType },
+                psiParentProducer = { containingElement.getParentOfType<UDeclaration>()?.psi ?: declaration.parent },
-                    ktInitializer = initializer,
+                containingElement = containingElement,
-                    psiParent = psiParent,
+                ktElement = declaration)
-                    containingElement = containingElement,
-                    ktElement = initializer)
-        }
 
-        fun create(name: String, localFunction: KtFunction, containingElement: UElement): PsiLocalVariable {
+        fun create(initializer: KtExpression, containingElement: UElement, parent: PsiElement): PsiLocalVariable = UastKotlinPsiVariable(
-            val psiParent = containingElement.getParentOfType<UDeclaration>()?.psi ?: localFunction.parent
+                manager = initializer.manager,
-            return UastKotlinPsiVariable(
+                name = "var" + Integer.toHexString(initializer.getHashCode()),
-                    manager = localFunction.manager,
+                typeProducer = { initializer.getType(containingElement) ?: UastErrorType },
-                    name = name,
+                ktInitializer = initializer,
-                    typeProducer = { localFunction.getFunctionType(containingElement) ?: UastErrorType },
+                psiParentProducer = { containingElement.getParentOfType<UDeclaration>()?.psi ?: parent },
-                    ktInitializer = localFunction,
+                containingElement = containingElement,
-                    psiParent = psiParent,
+                ktElement = initializer)
-                    containingElement = containingElement,
+
-                    ktElement = localFunction)
+        fun create(name: String, localFunction: KtFunction, containingElement: UElement): PsiLocalVariable = UastKotlinPsiVariable(
-        }
+                manager = localFunction.manager,
+                name = name,
+                typeProducer = { localFunction.getFunctionType(containingElement) ?: UastErrorType },
+                ktInitializer = localFunction,
+                psiParentProducer = { containingElement.getParentOfType<UDeclaration>()?.psi ?: localFunction.parent },
+                containingElement = containingElement,
+                ktElement = localFunction)
     }
 }
 

plugins/uast-kotlin/testData/WhenAndDestructing.kt

@@ -0,0 +1,17 @@
+fun getElementsAdditionalResolve(string: String): String {
+
+    val arr = listOf("1", "2")
+
+    when (string) {
+        "aaaa" -> {
+            return "bindingContext"
+        }
+
+        else -> {
+            val (bindingContext, statementFilter) = arr
+            return bindingContext
+        }
+    }
+}
+
+

plugins/uast-kotlin/testData/WhenAndDestructing.log.txt

@@ -0,0 +1,42 @@
+UFile (package = )
+    UClass (name = WhenAndDestructingKt)
+        UAnnotationMethod (name = getElementsAdditionalResolve)
+            UParameter (name = string)
+                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
+            UBlockExpression
+                UDeclarationsExpression
+                    ULocalVariable (name = arr)
+                        UCallExpression (kind = UastCallKind(name='method_call'), argCount = 2))
+                            UIdentifier (Identifier (listOf))
+                            USimpleNameReferenceExpression (identifier = listOf)
+                            ULiteralExpression (value = "1")
+                            ULiteralExpression (value = "2")
+                USwitchExpression
+                    USimpleNameReferenceExpression (identifier = string)
+                    UExpressionList (when)
+                        USwitchClauseExpressionWithBody
+                            ULiteralExpression (value = "aaaa")
+                            UExpressionList (when_entry)
+                                UReturnExpression
+                                    ULiteralExpression (value = "bindingContext")
+                                UBreakExpression (label = null)
+                        USwitchClauseExpressionWithBody
+                            UExpressionList (when_entry)
+                                UDeclarationsExpression
+                                    ULocalVariable (name = var837f1e82)
+                                        USimpleNameReferenceExpression (identifier = arr)
+                                    ULocalVariable (name = bindingContext)
+                                        UQualifiedReferenceExpression
+                                            USimpleNameReferenceExpression (identifier = var837f1e82)
+                                            UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                                                UIdentifier (Identifier (component1))
+                                                USimpleNameReferenceExpression (identifier = <anonymous class>)
+                                    ULocalVariable (name = statementFilter)
+                                        UQualifiedReferenceExpression
+                                            USimpleNameReferenceExpression (identifier = var837f1e82)
+                                            UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                                                UIdentifier (Identifier (component2))
+                                                USimpleNameReferenceExpression (identifier = <anonymous class>)
+                                UReturnExpression
+                                    USimpleNameReferenceExpression (identifier = bindingContext)
+                                UBreakExpression (label = null)

plugins/uast-kotlin/testData/WhenAndDestructing.render.txt

@@ -0,0 +1,21 @@
+public final class WhenAndDestructingKt {
+    public static final fun getElementsAdditionalResolve(string: java.lang.String) : java.lang.String {
+        var arr: java.util.List<? extends java.lang.String> = listOf("1", "2")
+        switch (string)  {
+            "aaaa" -> {
+                return "bindingContext"
+                break
+            }
+
+             -> {
+                var var837f1e82: <ErrorType> = arr
+                var bindingContext: java.lang.String = var837f1e82.<anonymous class>()
+                var statementFilter: java.lang.String = var837f1e82.<anonymous class>()
+                return bindingContext
+                break
+            }
+
+        }
+
+    }
+}

plugins/uast-kotlin/tests/KotlinUastApiTest.kt

@@ -131,23 +131,19 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
     fun testWhenStringLiteral() {
         doTest("WhenStringLiteral") { _, file ->
 
-            fun UFile.findULiteralExpressionFromPsi(refText: String): ULiteralExpression =
+            file.findFromPsi<ULiteralExpression>("abc").let { literalExpression ->
-                    this.psi.findElementAt(file.psi.text.indexOf(refText))!!
-                            .parentsWithSelf.mapNotNull { it.toUElementOfType<ULiteralExpression>() }.first()
-
-            file.findULiteralExpressionFromPsi("abc").let { literalExpression ->
                 val psi = literalExpression.psi!!
                 Assert.assertTrue(psi is KtLiteralStringTemplateEntry)
                 UsefulTestCase.assertInstanceOf(literalExpression.uastParent, USwitchClauseExpressionWithBody::class.java)
             }
 
-            file.findULiteralExpressionFromPsi("def").let { literalExpression ->
+            file.findFromPsi<ULiteralExpression>("def").let { literalExpression ->
                 val psi = literalExpression.psi!!
                 Assert.assertTrue(psi is KtLiteralStringTemplateEntry)
                 UsefulTestCase.assertInstanceOf(literalExpression.uastParent, USwitchClauseExpressionWithBody::class.java)
             }
 
-            file.findULiteralExpressionFromPsi("def1").let { literalExpression ->
+            file.findFromPsi<ULiteralExpression>("def1").let { literalExpression ->
                 val psi = literalExpression.psi!!
                 Assert.assertTrue(psi is KtLiteralStringTemplateEntry)
                 UsefulTestCase.assertInstanceOf(literalExpression.uastParent, UBlockExpression::class.java)
@@ -156,4 +152,26 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
 
         }
     }
+
+    @Test
+    fun testWhenAndDestructing() {
+        doTest("WhenAndDestructing") { _, file ->
+
+            file.findFromPsi<UExpression>("val (bindingContext, statementFilter) = arr").let { e ->
+                val uBlockExpression = e.getParentOfType<UBlockExpression>()
+                Assert.assertNotNull(uBlockExpression)
+                val uMethod = uBlockExpression!!.getParentOfType<UMethod>()
+                Assert.assertNotNull(uMethod)
+            }
+
+        }
+    }
+}
+
+
+
+inline fun <reified T : UElement> UFile.findFromPsi(refText: String): T =
+        (this.psi.findElementAt(this.psi.text.indexOf(refText)) ?: throw AssertionError("requested text '$refText' was not found in ${this.psi.name}"))
+                .parentsWithSelf.dropWhile { !it.text.contains(refText) }.mapNotNull { it.toUElementOfType<T>() }.first().also {
+            if (it.psi != null && it.psi?.text != refText) throw AssertionError("requested text '$refText' found as '${it.psi?.text}' in $it")
 }

plugins/uast-kotlin/tests/SimpleKotlinRenderLogTest.kt

@@ -48,4 +48,7 @@ class SimpleKotlinRenderLogTest : AbstractKotlinRenderLogTest() {
     @Test fun testUnexpectedContainer() = doTest("UnexpectedContainerException") { testName, file -> check(testName, file, false) }
 
     @Test fun testWhenStringLiteral() = doTest("WhenStringLiteral") { testName, file -> check(testName, file, false) }
+
+    @Test
+    fun testWhenAndDestructing() = doTest("WhenAndDestructing") { testName, file -> check(testName, file, false) }
 }