UAST properly handles string "when conditions" in lazy parent (KT-20709) (#1338)

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

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.expressions.KotlinUElvisExpression
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
@@ -59,14 +60,16 @@ abstract class KotlinAbstractUElement(private val givenParent: UElement?) : UEle
         if (parent is KtStringTemplateEntryWithExpression) {
             parent = parent.parent
         }
-        if ((parent is KtStringTemplateExpression && parent.entries.size == 1) ||
-            parent is KtWhenConditionWithExpression) {
+        if ((parent is KtStringTemplateExpression && parent.entries.size == 1)) {
+            parent = parent.parent
+        }
+        if (parent is KtWhenConditionWithExpression) {
             parent = parent.parent
         }
 
         val result = doConvertParent(this, parent)
         if (result == this) {
-            throw IllegalStateException("Loop in parent structure when converting a $psi")
+            throw IllegalStateException("Loop in parent structure when converting a $psi of type ${psi?.javaClass} with parent $parent of type ${parent?.javaClass} text: [${parent?.text}]")
         }
 
         return result
@@ -113,7 +116,7 @@ fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
         return result.initializingClass
     }
 
-    if (result is USwitchClauseExpressionWithBody && element.psi !is KtWhenCondition) {
+    if (result is USwitchClauseExpressionWithBody && !isInConditionBranch(element, result)) {
         return result.body
     }
 
@@ -132,6 +135,10 @@ fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
     return result
 }
 
+private fun isInConditionBranch(element: UElement, result: USwitchClauseExpressionWithBody) =
+        element.psi?.parentsWithSelf?.takeWhile { it !== result.psi }?.any { it is KtWhenCondition } ?: false
+
+
 private fun findAnnotationClassFromConstructorParameter(parameter: KtParameter): UClass? {
     val primaryConstructor = parameter.getStrictParentOfType<KtPrimaryConstructor>() ?: return null
     val containingClass = primaryConstructor.getContainingClassOrObject()

plugins/uast-kotlin/testData/WhenStringLiteral.kt

@@ -0,0 +1,12 @@
+val a = readLine()
+
+val b = when(a) {
+    "abc" -> 1
+    "def", "ghi" -> 2
+    else -> 3
+}
+
+when(a) {
+    "abc1" -> println(1)
+    "def1", "ghi1" -> println(2)
+}

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

@@ -0,0 +1,33 @@
+UFile (package = )
+    UClass (name = WhenStringLiteralKt)
+        UField (name = a)
+            UAnnotation (fqName = org.jetbrains.annotations.Nullable)
+            UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                UIdentifier (Identifier (readLine))
+                USimpleNameReferenceExpression (identifier = readLine)
+        UField (name = b)
+            UAnnotation (fqName = null)
+            USwitchExpression
+                USimpleNameReferenceExpression (identifier = a)
+                UExpressionList (when)
+                    USwitchClauseExpressionWithBody
+                        ULiteralExpression (value = "abc")
+                        UExpressionList (when_entry)
+                            ULiteralExpression (value = 1)
+                            UBreakExpression (label = null)
+                    USwitchClauseExpressionWithBody
+                        ULiteralExpression (value = "def")
+                        ULiteralExpression (value = "ghi")
+                        UExpressionList (when_entry)
+                            ULiteralExpression (value = 2)
+                            UBreakExpression (label = null)
+                    USwitchClauseExpressionWithBody
+                        UExpressionList (when_entry)
+                            ULiteralExpression (value = 3)
+                            UBreakExpression (label = null)
+        UAnnotationMethod (name = getA)
+        UAnnotationMethod (name = getB)
+        UAnnotationMethod (name = <no name provided>)
+            UBlockExpression
+                ULiteralExpression (value = "abc1")
+                ULiteralExpression (value = "def1")

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

@@ -0,0 +1,27 @@
+public final class WhenStringLiteralKt {
+    private static final var a: java.lang.String = readLine()
+    private static final var b: int = switch (a)  {
+        "abc" -> {
+            1
+            break
+        }
+        
+        "def", "ghi" -> {
+            2
+            break
+        }
+        
+         -> {
+            3
+            break
+        }
+        
+    }
+    
+    public static final fun getA() : java.lang.String = UastEmptyExpression
+    public static final fun getB() : int = UastEmptyExpression
+    public static final fun <no name provided>() : void {
+        "abc1"
+        "def1"
+    }
+}

plugins/uast-kotlin/tests/KotlinUastApiTest.kt

@@ -6,6 +6,7 @@ import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtLiteralStringTemplateEntry
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.test.testFramework.KtUsefulTestCase
 import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.KotlinUastLanguagePlugin
@@ -125,4 +126,34 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
             UsefulTestCase.assertInstanceOf(uBinaryExpression.uastParent, UIfExpression::class.java)
         }
     }
+
+    @Test
+    fun testWhenStringLiteral() {
+        doTest("WhenStringLiteral") { _, file ->
+
+            fun UFile.findULiteralExpressionFromPsi(refText: String): ULiteralExpression =
+                    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 ->
+                val psi = literalExpression.psi!!
+                Assert.assertTrue(psi is KtLiteralStringTemplateEntry)
+                UsefulTestCase.assertInstanceOf(literalExpression.uastParent, USwitchClauseExpressionWithBody::class.java)
+            }
+
+            file.findULiteralExpressionFromPsi("def1").let { literalExpression ->
+                val psi = literalExpression.psi!!
+                Assert.assertTrue(psi is KtLiteralStringTemplateEntry)
+                UsefulTestCase.assertInstanceOf(literalExpression.uastParent, UBlockExpression::class.java)
+            }
+
+
+        }
+    }
 }

plugins/uast-kotlin/tests/SimpleKotlinRenderLogTest.kt

@@ -46,4 +46,6 @@ class SimpleKotlinRenderLogTest : AbstractKotlinRenderLogTest() {
     @Test fun testParametersWithDefaultValues() = doTest("ParametersWithDefaultValues")
 
     @Test fun testUnexpectedContainer() = doTest("UnexpectedContainerException") { testName, file -> check(testName, file, false) }
+
+    @Test fun testWhenStringLiteral() = doTest("WhenStringLiteral") { testName, file -> check(testName, file, false) }
 }