[java-decompiler] IDEA-353923 Wrap return statements into braces in switch expression

GitOrigin-RevId: 169162fe9d0483a39b151cbbe4a54352cc78daa0

plugins/java-decompiler/engine/src/org/jetbrains/java/decompiler/modules/decompiler/stats/SwitchStatement.java

@@ -14,10 +14,7 @@ import org.jetbrains.java.decompiler.modules.decompiler.StatEdge;
 import org.jetbrains.java.decompiler.modules.decompiler.StatEdge.EdgeDirection;
 import org.jetbrains.java.decompiler.modules.decompiler.StatEdge.EdgeType;
 import org.jetbrains.java.decompiler.modules.decompiler.SwitchHelper;
-import org.jetbrains.java.decompiler.modules.decompiler.exps.ConstExprent;
-import org.jetbrains.java.decompiler.modules.decompiler.exps.Exprent;
-import org.jetbrains.java.decompiler.modules.decompiler.exps.FieldExprent;
-import org.jetbrains.java.decompiler.modules.decompiler.exps.SwitchExprent;
+import org.jetbrains.java.decompiler.modules.decompiler.exps.*;
 import org.jetbrains.java.decompiler.struct.gen.VarType;
 import org.jetbrains.java.decompiler.util.TextBuffer;
 
@@ -228,7 +225,21 @@ public final class SwitchStatement extends Statement {
         tracer.incrementCurrentSourceLine();
       }
       else {
-        buf.append(ExprProcessor.jmpWrapper(stat, canBeRule ? 0 : indent + 2, false, tracer));
+        if (canBeRule && stat instanceof BasicBlockStatement blockStatement && blockStatement.getExprents() != null &&
+            (blockStatement.getExprents().size() > 1 ||
+             (blockStatement.getExprents().size() == 1 &&
+              blockStatement.getExprents().get(0) instanceof ExitExprent exitExprent &&
+              exitExprent.getExitType() == ExitExprent.EXIT_RETURN))) {
+          TextBuffer buffer = ExprProcessor.jmpWrapper(stat, indent + 2, false, tracer);
+          buf.append("{").appendLineSeparator();
+          tracer.incrementCurrentSourceLine();
+          buf.append(buffer).appendIndent(indent + 1).append("}").appendLineSeparator();
+          tracer.incrementCurrentSourceLine();
+        }
+        else {
+          TextBuffer buffer = ExprProcessor.jmpWrapper(stat, canBeRule ? 0 : indent + 2, false, tracer);
+          buf.append(buffer);
+        }
       }
     }
     buf.appendIndent(indent).append("}").appendLineSeparator();

plugins/java-decompiler/engine/test/org/jetbrains/java/decompiler/SingleClassesTest.java

@@ -164,6 +164,7 @@ public class SingleClassesTest {
   @Test public void testComplexInstanceOfRecordPatternJavac() { doTest("pkg/TestComplexInstanceOfRecordPatternJavac"); }
   @Test public void testSwitchWithDeconstructionsWithoutNestedJavac() { doTest("pkg/TestSwitchWithDeconstructionsWithoutNestedJavac"); }
   @Test public void testSwitchNestedDeconstructionJavac() { doTest("pkg/TestSwitchNestedDeconstructionsJavac"); }
+  @Test public void testSwitchWrapReturnJavac() { doTest("pkg/TestSwitchWrapReturnJavac"); }
 
   // TODO: fix all below
   //@Test public void testUnionType() { doTest("pkg/TestUnionType"); }

plugins/java-decompiler/engine/testData/results/TestSwitchWrapReturnJavac.dec

@@ -0,0 +1,61 @@
+class TestSwitchWrapReturnJavac {
+   public static void main(String[] args) {
+      System.out.println(test(new S.A()));// 14
+   }// 15
+
+   private static int test(S a) {
+      switch (a) {// 18
+         case S.A var3 -> {// 20
+            return 1;
+         }
+         case S.B var4 -> {// 23
+            return 2;
+         }
+         case S.C var5 -> {// 26
+            return 3;
+         }
+      }
+   }
+
+   sealed interface S {
+      public static record C() implements S {
+      }
+
+      public static record B() implements S {
+      }
+
+      public static record A() implements S {
+      }
+   }
+}
+
+class 'TestSwitchWrapReturnJavac' {
+   method 'main ([Ljava/lang/String;)V' {
+      0      2
+      a      2
+      d      2
+      10      3
+   }
+
+   method 'test (LTestSwitchWrapReturnJavac$S;)I' {
+      10      6
+      3b      7
+      3c      7
+      43      10
+      44      10
+      4b      13
+      4c      13
+   }
+}
+
+Lines mapping:
+14 <-> 3
+15 <-> 4
+18 <-> 7
+20 <-> 8
+23 <-> 11
+26 <-> 14
+Not mapped:
+19
+22
+25

plugins/java-decompiler/engine/testData/src/pkg/TestSwitchWrapReturnJavac.java

@@ -0,0 +1,30 @@
+class TestSwitchWrapReturnJavac {
+  sealed interface S {
+    record A() implements S {
+    }
+
+    record B() implements S {
+    }
+
+    record C() implements S {
+    }
+  }
+
+  public static void main(String[] args) {
+    System.out.println(test(new S.A()));
+  }
+
+  private static int test(S a) {
+    switch (a) {
+      case S.A a1 -> {
+        return 1;
+      }
+      case S.B b -> {
+        return 2;
+      }
+      case S.C c -> {
+        return 3;
+      }
+    }
+  }
+}