[java-decompiler] IDEA-326015 Support common cases of deconstruction…

… in switch - collapse default case with assignment into default patterns GitOrigin-RevId: b9458df252c6cb6d1215cbbf695c510779e2816f
posledniypoet · Nov 17, 2023 · 5d10638 · 5d10638
1 parent ee778ce
commit 5d10638
Show file tree

Hide file tree

Showing 3 changed files with 259 additions and 214 deletions.
diff --git a/src/org/jetbrains/java/decompiler/modules/decompiler/SwitchPatternHelper.java b/src/org/jetbrains/java/decompiler/modules/decompiler/SwitchPatternHelper.java
@@ -847,7 +847,9 @@ private boolean edgeCanBeWayOutOfRoot(@Nullable DoStatement doStatement, @NotNul
       return edge.getDestination() instanceof DummyExitStatement ||
              (doStatement == null && edge.closure == myRootSwitchStatement) ||
              (doStatement != null && edge.closure == doStatement) ||
-             (doStatement != null && edge.closure instanceof DoStatement upperDoStatement && upperDoStatement.containsStatement(doStatement)) ||
+             (doStatement != null &&
+              edge.closure instanceof DoStatement upperDoStatement &&
+              upperDoStatement.containsStatement(doStatement)) ||
              outsideCatch(edge, doStatement);
     }
 
@@ -1165,6 +1167,7 @@ private static class SwitchOnReferenceCandidate implements SwitchOnCandidate {
     private final SwitchStatement myRootSwitchStatement;
     @NotNull
     private final InvocationExprent myPreviousSelector;
+    @NotNull
     private final Exprent myNewSwitchSelectorVariant;
     @NotNull
     private final List<TempVarAssignmentItem> myTempVarAssignments;
@@ -1225,12 +1228,80 @@ public void simplify() {
       if (headExprent != null) {
         headExprent.replaceExprent(myPreviousSelector, myNewSwitchSelectorVariant);
       }
-
       if (hasPattern) {
         remapWithPatterns(myRootSwitchStatement, myPatternContainer, myUppedDoStatement, myTempVarAssignments);
         cleanDefault(myRootSwitchStatement);
       }
+      Exprent oldSelector = myPreviousSelector.getInstance();
+      if (oldSelector instanceof VarExprent oldSelectorVarExprent &&
+          myNewSwitchSelectorVariant instanceof VarExprent newSwitchSelectorVarExprent) {
+        changeEverywhere(oldSelectorVarExprent, newSwitchSelectorVarExprent, myRootSwitchStatement.getCaseStatements());
+        changeDefaultToFullCase(myRootSwitchStatement, newSwitchSelectorVarExprent);
+      }
+    }
+
+    private static void changeDefaultToFullCase(@NotNull SwitchStatement myRoot, @NotNull VarExprent newSwitch) {
+      List<List<StatEdge>> edges = myRoot.getCaseEdges();
+      for (int i = 0; i < edges.size(); i++) {
+        List<StatEdge> statEdges = edges.get(i);
+        if (statEdges.size() == 1 && statEdges.get(0) == myRoot.getDefaultEdge()) {
+          Statement defaultStatement = myRoot.getCaseStatements().get(i);
+          Statement statementWithFirstAssignment = getStatementWithFirstAssignment(defaultStatement);
+          if (statementWithFirstAssignment!=null &&
+              statementWithFirstAssignment.getExprents() != null &&
+              !statementWithFirstAssignment.getExprents().isEmpty() &&
+              statementWithFirstAssignment.getExprents().get(0) instanceof AssignmentExprent assignmentExprent &&
+              assignmentExprent.getRight() != null &&
+              assignmentExprent.getLeft() instanceof VarExprent newVarExprent &&
+              assignmentExprent.getLeft().getExprType() != null &&
+              assignmentExprent.getRight().equals(newSwitch) &&
+              assignmentExprent.getLeft().getExprType().equals(newSwitch.getExprType())) {
+            statementWithFirstAssignment.getExprents().remove(0);
+            List<@Nullable Exprent> defaultValues = myRoot.getCaseValues().get(i);
+            defaultValues.clear();
+            defaultValues.add(newVarExprent);
+            myRoot.setUseCustomDefault();
+            break;
+          }
+        }
+      }
     }
+
+    @Nullable
+    private static Statement getStatementWithFirstAssignment(@NotNull Statement statement) {
+      if (statement.getExprents() != null && !statement.getExprents().isEmpty()) {
+        if (statement.getExprents().get(0) instanceof AssignmentExprent) {
+          return statement;
+        }
+        else {
+          return null;
+        }
+      }
+      if (!statement.getStats().isEmpty()) {
+        return getStatementWithFirstAssignment(statement.getStats().get(0));
+      }
+      return null;
+    }
+
+    private static void changeEverywhere(@NotNull VarExprent oldVariant,
+                                         @NotNull VarExprent newVariant,
+                                         @NotNull List<Statement> statements) {
+      for (Statement statement : statements) {
+        if (statement.getExprents() != null) {
+          for (Exprent exprent : statement.getExprents()) {
+            for (Exprent nestedExprent : exprent.getAllExprents()) {
+              if (nestedExprent.equals(oldVariant)) {
+                exprent.replaceExprent(nestedExprent, newVariant);
+              }
+            }
+          }
+        }
+        for (Statement nestedStat : statement.getStats()) {
+          changeEverywhere(oldVariant, newVariant, nestedStat.getStats());
+        }
+      }
+    }
+
     @Override
     public Set<SwitchStatement> usedSwitch() {
       return myUsedSwitchStatements;

diff --git a/src/org/jetbrains/java/decompiler/modules/decompiler/stats/SwitchStatement.java b/src/org/jetbrains/java/decompiler/modules/decompiler/stats/SwitchStatement.java
@@ -35,6 +35,7 @@ public final class SwitchStatement extends Statement {
   private StatEdge defaultEdge;
   private Exprent headExprent;
   private boolean canBeRule = false;
+  private boolean useCustomDefault = false;
 
   private SwitchStatement() {
     super(StatementType.SWITCH);
@@ -134,6 +135,10 @@ public int duplicateCaseStatement(@NotNull Statement currentStatement) {
     return caseIndex + 1;
   }
 
+  public void setUseCustomDefault() {
+    useCustomDefault = true;
+  }
+
   @Nullable
   public static Statement isHead(@NotNull Statement head) {
     if (head.type == StatementType.BASIC_BLOCK && head.getLastBasicType() == StatementType.SWITCH) {
@@ -172,7 +177,7 @@ public TextBuffer toJava(int indent, @NotNull BytecodeMappingTracer tracer) {
       List<StatEdge> edges = caseEdges.get(i);
       List<Exprent> values = caseValues.get(i);
       for (int j = 0; j < edges.size(); j++) {
-        if (edges.get(j) == defaultEdge) {
+        if (edges.get(j) == defaultEdge && !useCustomDefault) {
           if (!canBeRule) {
             buf.appendIndent(indent + 1).append("default:").appendLineSeparator();
           }