Uast: support for UYieldExpression (KT-35574)

2020-04-03 14:12:49 +03:00
parent 0a2429c6f6
commit 9674643254
16 changed files with 572 additions and 35 deletions
@@ -193,6 +193,10 @@ fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
    }

    if (result is USwitchClauseExpressionWithBody && !isInConditionBranch(element, result)) {
+        val uYieldExpression = result.body.expressions.lastOrNull().safeAs<UYieldExpression>()
+        if (uYieldExpression != null && uYieldExpression.expression == element)
+            return uYieldExpression
+
        return result.body
    }

@@ -0,0 +1,264 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.uast.kotlin
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiMethod
+import org.jetbrains.kotlin.asJava.LightClassUtil
+import org.jetbrains.kotlin.asJava.elements.KtLightElement
+import org.jetbrains.kotlin.asJava.elements.KtLightMethod
+import org.jetbrains.kotlin.asJava.toLightGetter
+import org.jetbrains.kotlin.asJava.toLightSetter
+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.isPropertyParameter
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import org.jetbrains.uast.*
+import org.jetbrains.uast.kotlin.expressions.KotlinLocalFunctionULambdaExpression
+import org.jetbrains.uast.kotlin.expressions.KotlinUElvisExpression
+import org.jetbrains.uast.kotlin.internal.KotlinUElementWithComments
+import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
+import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
+
+abstract class KotlinAbstractUElement(private val givenParent: UElement?) : KotlinUElementWithComments,
+    JvmDeclarationUElementPlaceholder {
+
+    final override val uastParent: UElement? by lz {
+        givenParent ?: convertParent()
+    }
+
+    protected open fun convertParent(): UElement? {
+        @Suppress("DEPRECATION")
+        val psi = psi //TODO: `psi` is deprecated but it seems that it couldn't be simply replaced for this case
+        var parent = psi?.parent ?: sourcePsi?.parent ?: psi?.containingFile
+
+        if (psi is PsiMethod && psi !is KtLightMethod) { // handling of synthetic things not represented in lightclasses directly
+            when (parent) {
+                is KtClassBody -> {
+                    val grandParent = parent.parent
+                    doConvertParent(this, grandParent)?.let { return it }
+                    parent = grandParent
+                }
+                is KtFile -> {
+                    parent.toUElementOfType<UClass>()?.let { return it } // mutlifile facade class
+                }
+            }
+
+        }
+
+        if (psi is KtLightElement<*, *> && sourcePsi.safeAs<KtClassOrObject>()?.isLocal == true) {
+            val originParent = psi.kotlinOrigin?.parent
+            parent = when (originParent) {
+                null -> parent
+                is KtClassBody -> originParent.parent
+                else -> originParent
+            }
+        }
+
+        if (psi is KtAnnotationEntry) {
+            val parentUnwrapped = KotlinConverter.unwrapElements(parent) ?: return null
+            val target = psi.useSiteTarget?.getAnnotationUseSiteTarget()
+            when (target) {
+                AnnotationUseSiteTarget.PROPERTY_GETTER ->
+                    parent = (parentUnwrapped as? KtProperty)?.getter
+                             ?: (parentUnwrapped as? KtParameter)?.toLightGetter()
+                             ?: parent
+
+                AnnotationUseSiteTarget.PROPERTY_SETTER ->
+                    parent = (parentUnwrapped as? KtProperty)?.setter
+                             ?: (parentUnwrapped as? KtParameter)?.toLightSetter()
+                             ?: parent
+                AnnotationUseSiteTarget.FIELD ->
+                    parent = (parentUnwrapped as? KtProperty)
+                        ?: (parentUnwrapped as? KtParameter)
+                            ?.takeIf { it.isPropertyParameter() }
+                            ?.let(LightClassUtil::getLightClassBackingField)
+                                ?: parent
+                AnnotationUseSiteTarget.SETTER_PARAMETER ->
+                    parent = (parentUnwrapped as? KtParameter)
+                        ?.toLightSetter()?.parameterList?.parameters?.firstOrNull() ?: parent
+            }
+        }
+        if ((psi is UastKotlinPsiVariable || psi is UastKotlinPsiParameter) && parent != null) {
+            parent = parent.parent
+        }
+
+        if (KotlinConverter.forceUInjectionHost) {
+            if (parent is KtBlockStringTemplateEntry) {
+                parent = parent.parent
+            }
+        } else
+            while (parent is KtStringTemplateEntryWithExpression || parent is KtStringTemplateExpression && parent.entries.size == 1) {
+                parent = parent.parent
+            }
+
+        if (parent is KtWhenConditionWithExpression) {
+            parent = parent.parent
+        }
+
+        if (parent is KtImportList) {
+            parent = parent.parent
+        }
+
+        if (psi is KtFunctionLiteral && parent is KtLambdaExpression) {
+            parent = parent.parent
+        }
+
+        if (parent is KtLambdaArgument) {
+            parent = parent.parent
+        }
+
+        if (psi is KtSuperTypeCallEntry) {
+            parent = parent?.parent
+        }
+
+        if (parent is KtPropertyDelegate) {
+            parent = parent.parent
+        }
+
+        val result = doConvertParent(this, parent)
+        if (result == this) {
+            throw IllegalStateException("Loop in parent structure when converting a $psi of type ${psi?.javaClass} with parent $parent of type ${parent?.javaClass} text: [${parent?.text}], result = $result")
+        }
+
+        return result
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (other !is UElement) {
+            return false
+        }
+
+        return this.psi == other.psi
+    }
+
+    override fun hashCode() = psi?.hashCode() ?: 0
+}
+
+fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
+    val parentUnwrapped = KotlinConverter.unwrapElements(parent) ?: return null
+    if (parent is KtValueArgument && parentUnwrapped is KtAnnotationEntry) {
+        return (KotlinUastLanguagePlugin().convertElementWithParent(parentUnwrapped, null) as? KotlinUAnnotation)
+            ?.findAttributeValueExpression(parent)
+    }
+
+    if (parent is KtParameter) {
+        val annotationClass = findAnnotationClassFromConstructorParameter(parent)
+        if (annotationClass != null) {
+            return annotationClass.methods.find { it.name == parent.name }
+        }
+    }
+
+    if (parent is KtClassInitializer) {
+        val containingClass = parent.containingClassOrObject
+        if (containingClass != null) {
+            val containingUClass = KotlinUastLanguagePlugin().convertElementWithParent(containingClass, null) as? KotlinUClass
+            containingUClass?.methods?.filterIsInstance<KotlinConstructorUMethod>()?.firstOrNull { it.isPrimary }?.let {
+                return it.uastBody
+            }
+        }
+    }
+
+    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
+    }
+
+    if (result is UCallExpression && result.uastParent is UEnumConstant) {
+        return result.uastParent
+    }
+
+    if (result is USwitchClauseExpressionWithBody && !isInConditionBranch(element, result)) {
+        return result.body
+    }
+
+    if (result is KotlinUDestructuringDeclarationExpression &&
+        when (parent) {
+            is KtDestructuringDeclaration -> parent.initializer?.let { it == element.psi } == true
+            is KtDeclarationModifierList -> parent == element.sourcePsi?.parent
+            else -> false
+        }
+    ) {
+        return result.tempVarAssignment
+    }
+
+    if (result is KotlinUElvisExpression && parentUnwrapped is KtBinaryExpression) {
+        val branch: Sequence<PsiElement?> = element.psi?.parentsWithSelf.orEmpty().takeWhile { it != parentUnwrapped }
+        if (branch.contains(parentUnwrapped.left))
+            return result.lhsDeclaration
+        if (branch.contains(parentUnwrapped.right))
+            return result.rhsIfExpression
+    }
+
+    if ((result is UMethod || result is KotlinLocalFunctionULambdaExpression)
+        && result !is KotlinConstructorUMethod // no sense to wrap super calls with `return`
+        && element is UExpression
+        && element !is UBlockExpression
+        && element !is UTypeReferenceExpression // when element is a type in extension methods
+    ) {
+        return KotlinUBlockExpression.KotlinLazyUBlockExpression(result, { block ->
+            listOf(KotlinUImplicitReturnExpression(block).apply { returnExpression = element })
+        }).expressions.single()
+    }
+
+    if (result is KotlinULambdaExpression.Body && element is UExpression && result.implicitReturn?.returnExpression == element) {
+        return result.implicitReturn!!
+    }
+
+    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()
+    if (containingClass.isAnnotation()) {
+        return KotlinUastLanguagePlugin().convertElementWithParent(containingClass, null) as? UClass
+    }
+    return null
+}
+
+abstract class KotlinAbstractUExpression(givenParent: UElement?) :
+    KotlinAbstractUElement(givenParent),
+    UExpression,
+    JvmDeclarationUElementPlaceholder {
+
+    override val javaPsi: PsiElement? = null
+
+    override val psi
+        get() = sourcePsi
+
+    override val annotations: List<UAnnotation>
+        get() {
+            val annotatedExpression = sourcePsi?.parent as? KtAnnotatedExpression ?: return emptyList()
+            return annotatedExpression.annotationEntries.map { KotlinUAnnotation(it, this) }
+        }
+}
+
@@ -18,8 +18,10 @@ package org.jetbrains.uast.kotlin

 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtWhenEntry
 import org.jetbrains.kotlin.psi.KtWhenExpression
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.declarations.KotlinUIdentifier
 import org.jetbrains.uast.kotlin.kinds.KotlinSpecialExpressionKinds
@@ -64,13 +66,13 @@ class KotlinUSwitchEntry(
                expressions.forEach { appendln(it.asRenderString().withMargin) }
                appendln("}")
            }
-        }.apply {
+        }.apply KotlinUExpressionList@{
            val exprPsi = this@KotlinUSwitchEntry.sourcePsi.expression
            val userExpressions = when (exprPsi) {
                is KtBlockExpression -> exprPsi.statements.map { KotlinConverter.convertOrEmpty(it, this) }
                else -> listOf(KotlinConverter.convertOrEmpty(exprPsi, this))
            }
-            expressions = userExpressions + object : UBreakExpression, JvmDeclarationUElementPlaceholder {
+            expressions = userExpressions.subList(0, userExpressions.lastIndex) + object : UYieldExpression, JvmDeclarationUElementPlaceholder {
                override val javaPsi: PsiElement? = null
                override val sourcePsi: PsiElement? = null
                override val psi: PsiElement?
@@ -78,9 +80,12 @@ class KotlinUSwitchEntry(
                override val label: String?
                    get() = null
                override val uastParent: UElement?
-                    get() = this@KotlinUSwitchEntry
+                    get() = this@KotlinUExpressionList
                override val annotations: List<UAnnotation>
                    get() = emptyList()
+                override val expression: UExpression?
+                    get() = userExpressions.lastOrNull()?.sourcePsi.safeAs<KtExpression>()
+                        ?.let { KotlinConverter.convertExpression(it, this, DEFAULT_EXPRESSION_TYPES_LIST) }
            }
        }
    }
@@ -0,0 +1,94 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.uast.kotlin
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtWhenEntry
+import org.jetbrains.kotlin.psi.KtWhenExpression
+import org.jetbrains.uast.*
+import org.jetbrains.uast.kotlin.declarations.KotlinUIdentifier
+import org.jetbrains.uast.kotlin.kinds.KotlinSpecialExpressionKinds
+
+class KotlinUSwitchExpression(
+        override val sourcePsi: KtWhenExpression,
+        givenParent: UElement?
+) : KotlinAbstractUExpression(givenParent), USwitchExpression, KotlinUElementWithType {
+    override val expression by lz { KotlinConverter.convertOrNull(sourcePsi.subjectExpression, this) }
+
+    override val body: UExpressionList by lz {
+        object : KotlinUExpressionList(sourcePsi, KotlinSpecialExpressionKinds.WHEN, this@KotlinUSwitchExpression) {
+            override fun asRenderString() = expressions.joinToString("\n") { it.asRenderString().withMargin }
+        }.apply {
+            expressions = this@KotlinUSwitchExpression.sourcePsi.entries.map { KotlinUSwitchEntry(it, this) }
+        }
+    }
+
+    override fun asRenderString() = buildString {
+        val expr = expression?.let { "(" + it.asRenderString() + ") " } ?: ""
+        appendln("switch $expr {")
+        appendln(body.asRenderString())
+        appendln("}")
+    }
+
+    override val switchIdentifier: UIdentifier
+        get() = KotlinUIdentifier(null, this)
+}
+
+class KotlinUSwitchEntry(
+        override val sourcePsi: KtWhenEntry,
+        givenParent: UElement?
+) : KotlinAbstractUExpression(givenParent), USwitchClauseExpressionWithBody {
+    override val caseValues by lz {
+        sourcePsi.conditions.map { KotlinConverter.convertWhenCondition(it, this, DEFAULT_EXPRESSION_TYPES_LIST) ?: UastEmptyExpression(null) }
+    }
+
+    override val body: UExpressionList by lz {
+        object : KotlinUExpressionList(sourcePsi, KotlinSpecialExpressionKinds.WHEN_ENTRY, this@KotlinUSwitchEntry) {
+            override fun asRenderString() = buildString {
+                appendln("{")
+                expressions.forEach { appendln(it.asRenderString().withMargin) }
+                appendln("}")
+            }
+        }.apply {
+            val exprPsi = this@KotlinUSwitchEntry.sourcePsi.expression
+            val userExpressions = when (exprPsi) {
+                is KtBlockExpression -> exprPsi.statements.map { KotlinConverter.convertOrEmpty(it, this) }
+                else -> listOf(KotlinConverter.convertOrEmpty(exprPsi, this))
+            }
+            expressions = userExpressions + object : UBreakExpression, JvmDeclarationUElementPlaceholder {
+                override val javaPsi: PsiElement? = null
+                override val sourcePsi: PsiElement? = null
+                override val psi: PsiElement?
+                    get() = null
+                override val label: String?
+                    get() = null
+                override val uastParent: UElement?
+                    get() = this@KotlinUSwitchEntry
+                override val annotations: List<UAnnotation>
+                    get() = emptyList()
+            }
+        }
+    }
+
+    override fun convertParent(): UElement? {
+        val result = KotlinConverter.unwrapElements(sourcePsi.parent)?.let { parentUnwrapped ->
+            KotlinUastLanguagePlugin().convertElementWithParent(parentUnwrapped, null)
+        }
+        return (result as? KotlinUSwitchExpression)?.body ?: result
+    }
+}
@@ -20,10 +20,10 @@ UFile (package = )
                            UPolyadicExpression (operator = +)
                                ULiteralExpression (value = "aaaa")
                            UExpressionList (when_entry)
-                                UReturnExpression
-                                    UPolyadicExpression (operator = +)
-                                        ULiteralExpression (value = "bindingContext")
-                                UBreakExpression (label = null)
+                                UYieldExpression
+                                    UReturnExpression
+                                        UPolyadicExpression (operator = +)
+                                            ULiteralExpression (value = "bindingContext")
                        USwitchClauseExpressionWithBody
                            UExpressionList (when_entry)
                                UDeclarationsExpression
@@ -44,6 +44,6 @@ UFile (package = )
                                            UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
                                                UIdentifier (Identifier (component2))
                                                USimpleNameReferenceExpression (identifier = <anonymous class>, resolvesTo = null)
-                                UReturnExpression
-                                    USimpleNameReferenceExpression (identifier = bindingContext)
-                                UBreakExpression (label = null)
+                                UYieldExpression
+                                    UReturnExpression
+                                        USimpleNameReferenceExpression (identifier = bindingContext)
@@ -0,0 +1,49 @@
+UFile (package = )
+    UClass (name = WhenAndDestructingKt)
+        UMethod (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, resolvesTo = null)
+                            UPolyadicExpression (operator = +)
+                                ULiteralExpression (value = "1")
+                            UPolyadicExpression (operator = +)
+                                ULiteralExpression (value = "2")
+                USwitchExpression
+                    USimpleNameReferenceExpression (identifier = string)
+                    UExpressionList (when)
+                        USwitchClauseExpressionWithBody
+                            UPolyadicExpression (operator = +)
+                                ULiteralExpression (value = "aaaa")
+                            UExpressionList (when_entry)
+                                UReturnExpression
+                                    UPolyadicExpression (operator = +)
+                                        ULiteralExpression (value = "bindingContext")
+                                UBreakExpression (label = null)
+                        USwitchClauseExpressionWithBody
+                            UExpressionList (when_entry)
+                                UDeclarationsExpression
+                                    ULocalVariable (name = var837f1e82)
+                                        UAnnotation (fqName = null)
+                                        USimpleNameReferenceExpression (identifier = arr)
+                                    ULocalVariable (name = bindingContext)
+                                        UAnnotation (fqName = null)
+                                        UQualifiedReferenceExpression
+                                            USimpleNameReferenceExpression (identifier = var837f1e82)
+                                            UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                                                UIdentifier (Identifier (component1))
+                                                USimpleNameReferenceExpression (identifier = <anonymous class>, resolvesTo = null)
+                                    ULocalVariable (name = statementFilter)
+                                        UAnnotation (fqName = null)
+                                        UQualifiedReferenceExpression
+                                            USimpleNameReferenceExpression (identifier = var837f1e82)
+                                            UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                                                UIdentifier (Identifier (component2))
+                                                USimpleNameReferenceExpression (identifier = <anonymous class>, resolvesTo = null)
+                                UReturnExpression
+                                    USimpleNameReferenceExpression (identifier = bindingContext)
+                                UBreakExpression (label = null)
@@ -3,16 +3,14 @@ public final class WhenAndDestructingKt {
        var arr: java.util.List<? extends java.lang.String> = listOf("1", "2")
        switch (string)  {
            "aaaa" -> {
-                return "bindingContext"
-                break
+                yield return "bindingContext"
            }

             -> {
                @null var var837f1e82: <ErrorType> = arr
                @null var bindingContext: java.lang.String = var837f1e82.<anonymous class>()
                @null var statementFilter: java.lang.String = var837f1e82.<anonymous class>()
-                return bindingContext
-                break
+                yield return bindingContext
            }

        }
@@ -0,0 +1,21 @@
+public final class WhenAndDestructingKt {
+    public static final fun getElementsAdditionalResolve(@org.jetbrains.annotations.NotNull 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
+            }
+
+             -> {
+                @null var var837f1e82: <ErrorType> = arr
+                @null var bindingContext: java.lang.String = var837f1e82.<anonymous class>()
+                @null var statementFilter: java.lang.String = var837f1e82.<anonymous class>()
+                return bindingContext
+                break
+            }
+
+        }
+
+    }
+}
@@ -13,13 +13,13 @@ UFile (package = )
                                    USimpleNameReferenceExpression (identifier = it)
                                    UTypeReferenceExpression (name = java.lang.String)
                                UExpressionList (when_entry)
-                                    USimpleNameReferenceExpression (identifier = bar)
-                                    UBreakExpression (label = null)
+                                    UYieldExpression
+                                        USimpleNameReferenceExpression (identifier = bar)
                            USwitchClauseExpressionWithBody
                                UBinaryExpressionWithType
                                    USimpleNameReferenceExpression (identifier = it)
                                    UTypeReferenceExpression (name = java.lang.String)
                                UExpressionList (when_entry)
-                                    UPolyadicExpression (operator = +)
-                                        ULiteralExpression (value = "<error>")
-                                    UBreakExpression (label = null)
+                                    UYieldExpression
+                                        UPolyadicExpression (operator = +)
+                                            ULiteralExpression (value = "<error>")
@@ -0,0 +1,25 @@
+UFile (package = )
+    UClass (name = WhenIsKt)
+        UMethod (name = foo)
+            UParameter (name = bar)
+                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
+            UBlockExpression
+                UReturnExpression
+                    USwitchExpression
+                        USimpleNameReferenceExpression (identifier = bar)
+                        UExpressionList (when)
+                            USwitchClauseExpressionWithBody
+                                UBinaryExpressionWithType
+                                    USimpleNameReferenceExpression (identifier = it)
+                                    UTypeReferenceExpression (name = java.lang.String)
+                                UExpressionList (when_entry)
+                                    USimpleNameReferenceExpression (identifier = bar)
+                                    UBreakExpression (label = null)
+                            USwitchClauseExpressionWithBody
+                                UBinaryExpressionWithType
+                                    USimpleNameReferenceExpression (identifier = it)
+                                    UTypeReferenceExpression (name = java.lang.String)
+                                UExpressionList (when_entry)
+                                    UPolyadicExpression (operator = +)
+                                        ULiteralExpression (value = "<error>")
+                                    UBreakExpression (label = null)
@@ -2,13 +2,11 @@ public final class WhenIsKt {
    public static final fun foo(@org.jetbrains.annotations.NotNull bar: java.lang.Object) : java.lang.String {
        return switch (bar)  {
            it is java.lang.String -> {
-                bar
-                break
+                yield bar
            }

            it !is java.lang.String -> {
-                "<error>"
-                break
+                yield "<error>"
            }

        }
@@ -0,0 +1,17 @@
+public final class WhenIsKt {
+    public static final fun foo(@org.jetbrains.annotations.NotNull bar: java.lang.Object) : java.lang.String {
+        return switch (bar)  {
+            it is java.lang.String -> {
+                bar
+                break
+            }
+
+            it !is java.lang.String -> {
+                "<error>"
+                break
+            }
+
+        }
+
+    }
+}
@@ -14,20 +14,20 @@ UFile (package = )
                        UPolyadicExpression (operator = +)
                            ULiteralExpression (value = "abc")
                        UExpressionList (when_entry)
-                            ULiteralExpression (value = 1)
-                            UBreakExpression (label = null)
+                            UYieldExpression
+                                ULiteralExpression (value = 1)
                    USwitchClauseExpressionWithBody
                        UPolyadicExpression (operator = +)
                            ULiteralExpression (value = "def")
                        UPolyadicExpression (operator = +)
                            ULiteralExpression (value = "ghi")
                        UExpressionList (when_entry)
-                            ULiteralExpression (value = 2)
-                            UBreakExpression (label = null)
+                            UYieldExpression
+                                ULiteralExpression (value = 2)
                    USwitchClauseExpressionWithBody
                        UExpressionList (when_entry)
-                            ULiteralExpression (value = 3)
-                            UBreakExpression (label = null)
+                            UYieldExpression
+                                ULiteralExpression (value = 3)
        UMethod (name = getA)
        UMethod (name = getB)
        UMethod (name = <no name provided>)
@@ -0,0 +1,38 @@
+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, resolvesTo = null)
+        UField (name = b)
+            UAnnotation (fqName = org.jetbrains.annotations.NotNull)
+            USwitchExpression
+                USimpleNameReferenceExpression (identifier = a)
+                UExpressionList (when)
+                    USwitchClauseExpressionWithBody
+                        UPolyadicExpression (operator = +)
+                            ULiteralExpression (value = "abc")
+                        UExpressionList (when_entry)
+                            ULiteralExpression (value = 1)
+                            UBreakExpression (label = null)
+                    USwitchClauseExpressionWithBody
+                        UPolyadicExpression (operator = +)
+                            ULiteralExpression (value = "def")
+                        UPolyadicExpression (operator = +)
+                            ULiteralExpression (value = "ghi")
+                        UExpressionList (when_entry)
+                            ULiteralExpression (value = 2)
+                            UBreakExpression (label = null)
+                    USwitchClauseExpressionWithBody
+                        UExpressionList (when_entry)
+                            ULiteralExpression (value = 3)
+                            UBreakExpression (label = null)
+        UMethod (name = getA)
+        UMethod (name = getB)
+        UMethod (name = <no name provided>)
+            UBlockExpression
+                UPolyadicExpression (operator = +)
+                    ULiteralExpression (value = "abc1")
+                UPolyadicExpression (operator = +)
+                    ULiteralExpression (value = "def1")
@@ -2,18 +2,15 @@ public final class WhenStringLiteralKt {
    @org.jetbrains.annotations.Nullable private static final var a: java.lang.String = readLine()
    @org.jetbrains.annotations.NotNull private static final var b: int = switch (a)  {
        "abc" -> {
-            1
-            break
+            yield 1
        }

        "def", "ghi" -> {
-            2
-            break
+            yield 2
        }

         -> {
-            3
-            break
+            yield 3
        }

    }
@@ -0,0 +1,27 @@
+public final class WhenStringLiteralKt {
+    @org.jetbrains.annotations.Nullable private static final var a: java.lang.String = readLine()
+    @org.jetbrains.annotations.NotNull 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"
+    }
+}