[uast-java] properly handles PsiBlockStatement inside PsiSwitchLabeledRuleStatement (IDEA-202555, IDEA-CR-40758)

uast/uast-java/src/org/jetbrains/uast/java/JavaAbstractUElement.kt

@@ -66,17 +66,26 @@ abstract class JavaAbstractUElement(givenParent: UElement?) : JavaUElementWithCo
 
 private fun JavaAbstractUElement.unwrapSwitch(uParent: UElement): UElement {
   when (uParent) {
-    is JavaUCodeBlockExpression -> {
+    is UBreakExpression -> return uParent.uastParent ?: uParent
+    is UBlockExpression -> {
       val codeBlockParent = uParent.uastParent
-      if (codeBlockParent is JavaUSwitchEntryList) {
-        if (branchHasElement(psi, codeBlockParent.psi) { it is PsiSwitchLabelStatementBase }) {
-          return codeBlockParent
+      when (codeBlockParent) {
+
+        is JavaUBlockExpression -> {
+          val sourcePsi = codeBlockParent.sourcePsi
+          if (sourcePsi is PsiBlockStatement && sourcePsi.parent is PsiSwitchLabeledRuleStatement)
+            (codeBlockParent.uastParent as? JavaUSwitchEntry)?.let { return it.body }
         }
-        val psiElement = psi ?: return uParent
-        return codeBlockParent.findUSwitchEntryForBodyStatementMember(psiElement)?.body ?: return codeBlockParent
-      }
-      if (codeBlockParent is JavaUSwitchExpression) {
-        return unwrapSwitch(codeBlockParent)
+
+        is JavaUSwitchEntryList -> {
+          if (branchHasElement(psi, codeBlockParent.psi) { it is PsiSwitchLabelStatementBase }) {
+            return codeBlockParent
+          }
+          val psiElement = psi ?: return uParent
+          return codeBlockParent.findUSwitchEntryForBodyStatementMember(psiElement)?.body ?: return codeBlockParent
+        }
+
+        is JavaUSwitchExpression -> return unwrapSwitch(codeBlockParent)
       }
       return uParent
     }

uast/uast-java/src/org/jetbrains/uast/java/controlStructures/JavaUSwitchExpression.kt

@@ -51,7 +51,14 @@ class JavaUSwitchEntryList(override val psi: PsiSwitchBlock, override val uastPa
     val result = mutableListOf<JavaUSwitchEntry>()
     for (statement in statements) {
       if (statement is PsiSwitchLabeledRuleStatement) {
-        result += JavaUSwitchEntry(listOf(statement), listOfNotNull(statement.body), this)
+        val body = statement.body
+        result += when (body) {
+          is PsiBlockStatement ->
+            JavaUSwitchEntry(listOf(statement), body.codeBlock.statements.toList(), this)
+          else ->
+            JavaUSwitchEntry(listOf(statement), listOfNotNull(body), this)
+        }
+
       }
       if (statement is PsiSwitchLabelStatement) {
         if (currentBody.isEmpty()) {

uast/uast-tests/java/Simple/EnhancedSwitch.java

@@ -5,7 +5,10 @@ public class Main {
 
     final String numericString =
       switch (str) {
-        case "foo" -> "FOO";
+        case "foo" -> {
+          System.out.println("here");
+          break "FOO";
+        }
         case "bar" -> "BAR";
         case "baz" -> "bAz";
         default -> "default";
@@ -18,6 +21,7 @@ public class Main {
         case "bar":
           break "BAR";
         case "baz":
+          System.out.println("here");
           break "bAz";
         default:
           break "default";

uast/uast-tests/java/Simple/EnhancedSwitch.log.txt

@@ -13,7 +13,14 @@ UFile (package = )
                                 USwitchClauseExpressionWithBody
                                     ULiteralExpression (value = "foo")
                                     UExpressionList (switch_entry)
-                                        ULiteralExpression (value = "FOO")
+                                        UQualifiedReferenceExpression
+                                            UQualifiedReferenceExpression
+                                                USimpleNameReferenceExpression (identifier = System)
+                                                USimpleNameReferenceExpression (identifier = out)
+                                            UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
+                                                UIdentifier (Identifier (println))
+                                                ULiteralExpression (value = "here")
+                                        UBreakExpression (label = null)
                                 USwitchClauseExpressionWithBody
                                     ULiteralExpression (value = "bar")
                                     UExpressionList (switch_entry)
@@ -42,6 +49,13 @@ UFile (package = )
                                 USwitchClauseExpressionWithBody
                                     ULiteralExpression (value = "baz")
                                     UExpressionList (switch_entry)
+                                        UQualifiedReferenceExpression
+                                            UQualifiedReferenceExpression
+                                                USimpleNameReferenceExpression (identifier = System)
+                                                USimpleNameReferenceExpression (identifier = out)
+                                            UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
+                                                UIdentifier (Identifier (println))
+                                                ULiteralExpression (value = "here")
                                         UBreakExpression (label = null)
                                 USwitchClauseExpressionWithBody
                                     UDefaultCaseExpression

uast/uast-tests/java/Simple/EnhancedSwitch.render.txt

@@ -3,7 +3,8 @@ public class Main {
         var str: var = "abc"
         final var numericString: java.lang.String = switch (str)
             "foo" -> {
-                "FOO"
+                System.out.println("here")
+                break
             }
 
             "bar" -> {
@@ -29,6 +30,7 @@ public class Main {
             }
 
             "baz" -> {
+                System.out.println("here")
                 break
             }