From 0ae509504c81f77d95cdf1dc0c474177cab538a0 Mon Sep 17 00:00:00 2001
From: Mike Strobel <mike.strobel@gmail.com>
Date: Fri, 18 Jun 2021 20:36:04 -0400
Subject: [PATCH] Decent support for switch expressions, but more comprehensive
 support would require rewriting them in the bytecode AST layer.

---
 .idea/inspectionProfiles/Project_Default.xml  |   3 +
 .../com/strobel/decompiler/ITextOutput.java   |   2 +
 .../strobel/decompiler/PlainTextOutput.java   |  20 +-
 .../languages/java/JavaOutputVisitor.java     | 156 +++++++-
 .../languages/java/ast/BreakStatement.java    |  36 +-
 .../java/ast/DepthFirstAstVisitor.java        |  10 +
 .../languages/java/ast/IAstVisitor.java       |   2 +
 .../languages/java/ast/JavaNameResolver.java  |  10 +
 .../languages/java/ast/SwitchExpression.java  |  71 ++++
 .../java/ast/SwitchExpressionArm.java         |  91 +++++
 .../languages/java/ast/SwitchSection.java     |   4 +-
 .../RewriteSwitchExpressionsTransform.java    | 374 ++++++++++++++++++
 .../transforms/TransformationPipeline.java    |   1 +
 13 files changed, 752 insertions(+), 28 deletions(-)
 create mode 100644 Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchExpression.java
 create mode 100644 Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchExpressionArm.java
 create mode 100644 Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/transforms/RewriteSwitchExpressionsTransform.java

diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml
index fe325ec4..5d9df32e 100644
--- a/.idea/inspectionProfiles/Project_Default.xml
+++ b/.idea/inspectionProfiles/Project_Default.xml
@@ -6,6 +6,9 @@
       <option name="REPORT_INDEXED_LOOP" value="false" />
       <option name="ignoreUntypedCollections" value="false" />
     </inspection_tool>
+    <inspection_tool class="IfStatementWithIdenticalBranches" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <option name="myHighlightWhenLastStatementIsCall" value="false" />
+    </inspection_tool>
     <inspection_tool class="LocalCanBeFinal" enabled="true" level="WARNING" enabled_by_default="true">
       <option name="REPORT_VARIABLES" value="true" />
       <option name="REPORT_PARAMETERS" value="true" />
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/ITextOutput.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/ITextOutput.java
index 23b7d4b4..480393e6 100644
--- a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/ITextOutput.java
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/ITextOutput.java
@@ -26,6 +26,8 @@ public interface ITextOutput {
     void indent();
     void unindent();
 
+    int indentDepth();
+
     void write(final char ch);
     void write(final String text);
 
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/PlainTextOutput.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/PlainTextOutput.java
index 1b5e501c..7766dc30 100644
--- a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/PlainTextOutput.java
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/PlainTextOutput.java
@@ -24,6 +24,8 @@
 import java.io.Writer;
 import java.lang.reflect.UndeclaredThrowableException;
 
+import static com.strobel.core.Comparer.coalesce;
+
 public class PlainTextOutput implements ITextOutput {
     private final static String NULL_TEXT = String.valueOf((Object) null);
 
@@ -73,7 +75,7 @@ protected void writeIndent() {
                 try {
                     _writer.write(indentToken);
                 }
-                catch (IOException e) {
+                catch (final IOException e) {
                     throw new UndeclaredThrowableException(e);
                 }
             }
@@ -97,6 +99,11 @@ public void indent() {
         ++_indent;
     }
 
+    @Override
+    public int indentDepth() {
+        return _indent;
+    }
+
     @Override
     public void unindent() {
         --_indent;
@@ -114,7 +121,7 @@ public void write(final char ch) {
             }
             column++;
         }
-        catch (IOException e) {
+        catch (final IOException e) {
             throw new UndeclaredThrowableException(e);
         }
     }
@@ -134,9 +141,10 @@ protected void writeRaw(final String text) {
         writeIndent();
 
         try {
-            final int length = text != null ? text.length() : NULL_TEXT.length();
+            final String effectiveText = coalesce(text, NULL_TEXT);
+            final int length = effectiveText.length();
 
-            _writer.write(text);
+            _writer.write(effectiveText);
 
             column += length;
 
@@ -157,7 +165,7 @@ else if (newLineSeen) {
                 }
             }
         }
-        catch (IOException e) {
+        catch (final IOException e) {
             throw new UndeclaredThrowableException(e);
         }
     }
@@ -215,7 +223,7 @@ public void writeLine() {
         try {
             _writer.write("\n");
         }
-        catch (IOException e) {
+        catch (final IOException e) {
             throw new UndeclaredThrowableException(e);
         }
         _needsIndent = true;
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/JavaOutputVisitor.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/JavaOutputVisitor.java
index 3fa0b4e3..366efc2f 100644
--- a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/JavaOutputVisitor.java
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/JavaOutputVisitor.java
@@ -36,7 +36,6 @@
 import com.strobel.decompiler.languages.java.ast.WildcardType;
 import com.strobel.decompiler.languages.java.utilities.TypeUtilities;
 import com.strobel.decompiler.patterns.*;
-import com.strobel.util.ContractUtils;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -49,6 +48,8 @@
 
 @SuppressWarnings("ConstantConditions")
 public final class JavaOutputVisitor implements IAstVisitor<Void, Void> {
+    private final static int MAX_SWITCH_EXPRESSION_ARM_VALUES_LENGTH = 40;
+
     final TextOutputFormatter formatter;
     final DecompilerSettings settings;
     final JavaFormattingOptions policy;
@@ -215,7 +216,7 @@ void closeBrace(final BraceStyle style) {
         formatter.closeBrace(style);
         lastWritten = LastWritten.Other;
     }
-    
+
     void writeIdentifier(final Identifier identifier, final String text) {
         // for now, set the start location to *here*
         identifier.setStartLocation(new TextLocation(output.getRow(), output.getColumn()));
@@ -236,7 +237,7 @@ boolean writeIdentifier(final String identifier, final Role<Identifier> identifi
         writeSpecialsUpToRole(identifierRole != null ? identifierRole : Roles.IDENTIFIER);
 
         boolean wroteSpace = false;
-        
+
         if (isKeyword(identifier, containerStack.peek())) {
             if (lastWritten == LastWritten.KeywordOrIdentifier) {
                 space();
@@ -259,7 +260,7 @@ else if (lastWritten == LastWritten.KeywordOrIdentifier) {
         }
 
         lastWritten = LastWritten.KeywordOrIdentifier;
-        
+
         return wroteSpace;
     }
 
@@ -646,7 +647,6 @@ else if (childNode instanceof ParameterReferenceNode) {
         }
         else if (childNode instanceof AstTypeMatch) {
             visitAstTypeMatch((AstTypeMatch) childNode);
-
         }
         else {
             writePrimitiveValue(childNode);
@@ -813,7 +813,9 @@ public Void visitJavaTokenNode(final JavaTokenNode node, final Void ignored) {
             endNode(modifierToken);
         }
         else {
-            throw ContractUtils.unsupported();
+            startNode(node);
+            writeKeyword(node.getText());
+            endNode(node);
         }
         return null;
     }
@@ -974,12 +976,32 @@ public Void visitExpressionStatement(final ExpressionStatement node, final Void
     @Override
     public Void visitBreakStatement(final BreakStatement node, final Void ignored) {
         startNode(node);
-        writeKeyword("break");
 
-        final String label = node.getLabel();
+        JavaTokenNode yt = node.getYieldToken();
 
-        if (!StringUtilities.isNullOrEmpty(label)) {
-            writeIdentifier(label, Roles.LABEL);
+        if (yt.isNull()) {
+            yt = node.getBreakToken();
+        }
+
+        if (yt.isNull()) {
+            writeKeyword("break");
+        }
+        else {
+            yt.acceptVisitor(this, ignored);
+        }
+
+        final Expression value = node.getValue();
+
+        if (value.isNull()) {
+            final String label = node.getLabel();
+
+            if (!StringUtilities.isNullOrEmpty(label)) {
+                writeIdentifier(label, Roles.LABEL);
+            }
+        }
+        else {
+            space();
+            value.acceptVisitor(this, ignored);
         }
 
         semicolon();
@@ -1165,8 +1187,109 @@ public Void visitSwitchSection(final SwitchSection node, final Void ignored) {
             first = false;
         }
 
-        final boolean isBlock = node.getStatements().size() == 1 &&
-                                firstOrDefault(node.getStatements()) instanceof BlockStatement;
+        writeSwitchSectionStatements(node.getStatements(), false);
+        endNode(node);
+        return null;
+    }
+
+    @Override
+    public Void visitSwitchExpression(final SwitchExpression node, final Void ignored) {
+        startNode(node);
+        writeKeyword(SwitchStatement.SWITCH_KEYWORD_ROLE);
+        space(policy.SpaceBeforeSwitchParentheses);
+        leftParenthesis();
+        space(policy.SpacesWithinSwitchParentheses);
+        node.getGoverningExpression().acceptVisitor(this, ignored);
+        space(policy.SpacesWithinSwitchParentheses);
+        rightParenthesis();
+        openBrace(policy.StatementBraceStyle);
+
+        if (policy.IndentSwitchBody) {
+            formatter.indent();
+        }
+
+        for (final SwitchExpressionArm section : node.getArms()) {
+            section.acceptVisitor(this, ignored);
+        }
+
+        if (policy.IndentSwitchBody) {
+            formatter.unindent();
+        }
+
+        closeBrace(policy.StatementBraceStyle);
+        endNode(node);
+
+        return null;
+    }
+
+    @Override
+    public Void visitSwitchExpressionArm(final SwitchExpressionArm node, final Void ignored) {
+        startNode(node);
+
+        final AstNodeCollection<Expression> values = node.getValues();
+        final boolean isDefault = values.isEmpty();
+
+        writeKeyword(isDefault ? SwitchExpressionArm.DEFAULT_KEYWORD_ROLE
+                               : SwitchExpressionArm.CASE_KEYWORD_ROLE);
+
+        if (!isDefault) {
+            final int startColumn = output.getIndentToken().length() * output.indentDepth();
+            final int endColumn = output.getColumn();
+
+            String padding = null;
+            boolean comma = false;
+
+            for (final Expression value : values) {
+                if (comma) {
+                    comma(value);
+                }
+
+                space();
+
+                if (output.getColumn() - endColumn >= MAX_SWITCH_EXPRESSION_ARM_VALUES_LENGTH) {
+                    newLine();
+                    if (padding == null) {
+                        padding = StringUtilities.repeat(' ', endColumn - startColumn);
+                    }
+                    output.write(padding);
+                }
+
+                value.acceptVisitor(this, ignored);
+                comma = true;
+            }
+        }
+
+        if (node.isClassicStyle()) {
+            writeToken(SwitchExpressionArm.COLON_ROLE);
+        }
+        else {
+            space();
+            writeToken(SwitchExpressionArm.ARROW_ROLE);
+        }
+
+        space();
+
+        final Statement first;
+        final AstNodeCollection<Statement> statements = node.getStatements();
+
+        if (statements.size() == 1 && !((first = statements.firstOrNullObject()) instanceof BlockStatement)) {
+            first.acceptVisitor(this, ignored);
+        }
+        else {
+            writeSwitchSectionStatements(statements, true);
+        }
+
+        endNode(node);
+        return null;
+    }
+
+    private void writeSwitchSectionStatements(final AstNodeCollection<Statement> statements, final boolean forceBraces) {
+        final boolean isBlock = statements.size() == 1 &&
+                                firstOrDefault(statements) instanceof BlockStatement;
+
+        if (forceBraces && !isBlock) {
+            openBrace(BraceStyle.EndOfLine);
+        }
 
         if (policy.IndentCaseBody && !isBlock) {
             formatter.indent();
@@ -1176,16 +1299,17 @@ public Void visitSwitchSection(final SwitchSection node, final Void ignored) {
             newLine();
         }
 
-        for (final Statement statement : node.getStatements()) {
-            statement.acceptVisitor(this, ignored);
+        for (final Statement statement : statements) {
+            statement.acceptVisitor(this, null);
         }
 
         if (policy.IndentCaseBody && !isBlock) {
             formatter.unindent();
         }
 
-        endNode(node);
-        return null;
+        if (forceBraces && !isBlock) {
+            closeBrace(BraceStyle.EndOfLine);
+        }
     }
 
     @Override
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/BreakStatement.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/BreakStatement.java
index b6f713fa..1a84ab11 100644
--- a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/BreakStatement.java
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/BreakStatement.java
@@ -17,25 +17,53 @@
 package com.strobel.decompiler.languages.java.ast;
 
 import com.strobel.core.StringUtilities;
+import com.strobel.decompiler.languages.TextLocation;
 import com.strobel.decompiler.patterns.INode;
 import com.strobel.decompiler.patterns.Match;
 
 public class BreakStatement extends Statement {
     public final static TokenRole BREAK_KEYWORD_ROLE = new TokenRole("break", TokenRole.FLAG_KEYWORD);
+    public final static TokenRole YIELD_KEYWORD_ROLE = new TokenRole("yield", TokenRole.FLAG_KEYWORD);
 
-    public BreakStatement( int offset) {
-        super( offset);
+    public BreakStatement(final int offset) {
+        super(offset);
     }
 
-    public BreakStatement( int offset, final String label) {
-        super( offset);
+    public BreakStatement(final int offset, final String label) {
+        super(offset);
         setLabel(label);
+        setChildByRole(BREAK_KEYWORD_ROLE, new JavaTokenNode(TextLocation.EMPTY));
     }
 
     public final JavaTokenNode getBreakToken() {
         return getChildByRole(BREAK_KEYWORD_ROLE);
     }
 
+    public final JavaTokenNode getYieldToken() {
+        return getChildByRole(YIELD_KEYWORD_ROLE);
+    }
+
+    public final Expression getValue() {
+        return getChildByRole(Roles.EXPRESSION);
+    }
+
+    public final void setValue(final Expression value) {
+        setChildByRole(Roles.EXPRESSION, value);
+    }
+
+    public final void setYield(final boolean isYield) {
+        if (isYield) {
+            final JavaTokenNode old = getChildByRole(BREAK_KEYWORD_ROLE);
+            old.remove();
+            setChildByRole(YIELD_KEYWORD_ROLE, old.isNull() ? new JavaTokenNode(TextLocation.EMPTY) : old);
+        }
+        else {
+            final JavaTokenNode old = getChildByRole(YIELD_KEYWORD_ROLE);
+            old.remove();
+            setChildByRole(BREAK_KEYWORD_ROLE, old.isNull() ? new JavaTokenNode(TextLocation.EMPTY) : old);
+        }
+    }
+
     public final JavaTokenNode getSemicolonToken() {
         return getChildByRole(Roles.SEMICOLON);
     }
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/DepthFirstAstVisitor.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/DepthFirstAstVisitor.java
index d42dc2bd..ee447eee 100644
--- a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/DepthFirstAstVisitor.java
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/DepthFirstAstVisitor.java
@@ -159,6 +159,16 @@ public S visitSwitchSection(final SwitchSection node, final T data) {
         return visitChildren(node, data);
     }
 
+    @Override
+    public S visitSwitchExpression(final SwitchExpression node, final T data) {
+        return visitChildren(node, data);
+    }
+
+    @Override
+    public S visitSwitchExpressionArm(final SwitchExpressionArm node, final T data) {
+        return visitChildren(node, data);
+    }
+
     @Override
     public S visitCaseLabel(final CaseLabel node, final T data) {
         return visitChildren(node, data);
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/IAstVisitor.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/IAstVisitor.java
index 16b277e7..a168d17a 100644
--- a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/IAstVisitor.java
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/IAstVisitor.java
@@ -42,6 +42,8 @@ public interface IAstVisitor<T, R> {
     R visitReturnStatement(ReturnStatement node, T data);
     R visitSwitchStatement(SwitchStatement node, T data);
     R visitSwitchSection(SwitchSection node, T data);
+    R visitSwitchExpression(SwitchExpression node, T data);
+    R visitSwitchExpressionArm(SwitchExpressionArm node, T data);
     R visitCaseLabel(CaseLabel node, T data);
     R visitThrowStatement(ThrowStatement node, T data);
     R visitCatchClause(CatchClause node, T data);
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/JavaNameResolver.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/JavaNameResolver.java
index 62423ab3..82878416 100644
--- a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/JavaNameResolver.java
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/JavaNameResolver.java
@@ -482,6 +482,16 @@ public Set<Object> visitSwitchSection(final SwitchSection node, final String nam
             return Collections.emptySet();
         }
 
+        @Override
+        public Set<Object> visitSwitchExpression(final SwitchExpression node, final String data) {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public Set<Object> visitSwitchExpressionArm(final SwitchExpressionArm node, final String data) {
+            return Collections.emptySet();
+        }
+
         @Override
         public Set<Object> visitCaseLabel(final CaseLabel node, final String name) {
             return Collections.emptySet();
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchExpression.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchExpression.java
new file mode 100644
index 00000000..ec9b662f
--- /dev/null
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchExpression.java
@@ -0,0 +1,71 @@
+package com.strobel.decompiler.languages.java.ast;
+
+import com.strobel.decompiler.patterns.INode;
+import com.strobel.decompiler.patterns.Match;
+import com.strobel.decompiler.patterns.Role;
+
+public class SwitchExpression extends Expression{
+    public final static SwitchExpression NULL = new NullSwitchExpression();
+
+    public final static Role<SwitchExpressionArm> ARM_ROLE = new Role<>("SwitchExpressionArm", SwitchExpressionArm.class, SwitchExpressionArm.NULL);
+    public final static Role<Expression> GOVERNING_EXPRESSION_ROLE = new Role<>("GoverningExpression", Expression.class, Expression.NULL);
+
+    public SwitchExpression() {
+        super(MYSTERY_OFFSET);
+    }
+
+    public SwitchExpression(final int offset) {
+        super(offset);
+    }
+    
+    public final Expression getGoverningExpression() {
+        return getChildByRole(GOVERNING_EXPRESSION_ROLE);
+    }
+
+    public final void setGoverningExpression(final Expression value) {
+        setChildByRole(GOVERNING_EXPRESSION_ROLE, value);
+    }
+
+    public final AstNodeCollection<SwitchExpressionArm> getArms() {
+        return getChildrenByRole(ARM_ROLE);
+    }
+
+    @Override
+    public <T, R> R acceptVisitor(final IAstVisitor<? super T, ? extends R> visitor, final T data) {
+        return visitor.visitSwitchExpression(this, data);
+    }
+
+    @Override
+    public boolean matches(final INode other, final Match match) {
+        if (other instanceof SwitchExpression) {
+            final SwitchExpression otherStatement = (SwitchExpression) other;
+
+            return !otherStatement.isNull() &&
+                   getGoverningExpression().matches(otherStatement.getGoverningExpression(), match) &&
+                   getArms().matches(otherStatement.getArms(), match);
+        }
+
+        return false;
+    }
+
+    private static final class NullSwitchExpression extends SwitchExpression {
+        public NullSwitchExpression() {
+            super(Expression.MYSTERY_OFFSET);
+        }
+
+        @Override
+        public final boolean isNull() {
+            return true;
+        }
+
+        @Override
+        public <T, R> R acceptVisitor(final IAstVisitor<? super T, ? extends R> visitor, final T data) {
+            return null;
+        }
+
+        @Override
+        public boolean matches(final INode other, final Match match) {
+            return other == null || other.isNull();
+        }
+    }
+}
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchExpressionArm.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchExpressionArm.java
new file mode 100644
index 00000000..66910831
--- /dev/null
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchExpressionArm.java
@@ -0,0 +1,91 @@
+package com.strobel.decompiler.languages.java.ast;
+
+import com.strobel.decompiler.patterns.INode;
+import com.strobel.decompiler.patterns.Match;
+import com.strobel.decompiler.patterns.Role;
+
+public class SwitchExpressionArm extends AstNode {
+    public final static SwitchExpressionArm NULL = new NullSwitchExpressionArm();
+
+    public final static Role<Expression> VALUE_ROLE = new Role<>("Value", Expression.class, Expression.NULL);
+    public final static TokenRole CASE_KEYWORD_ROLE = CaseLabel.CASE_KEYWORD_ROLE;
+    public final static TokenRole DEFAULT_KEYWORD_ROLE = CaseLabel.DEFAULT_KEYWORD_ROLE;
+    public final static TokenRole ARROW_ROLE = new TokenRole("->", TokenRole.FLAG_OPERATOR);
+    public final static TokenRole COLON_ROLE = new TokenRole(":", TokenRole.FLAG_DELIMITER);
+
+    public final AstNodeCollection<Statement> getStatements() {
+        return getChildrenByRole(Roles.EMBEDDED_STATEMENT);
+    }
+
+    public final AstNodeCollection<Expression> getValues() {
+        return getChildrenByRole(VALUE_ROLE);
+    }
+
+    private boolean defaultCase;
+    private boolean isClassicStyle;
+
+    public boolean isDefaultCase() {
+        return defaultCase;
+    }
+
+    public void setDefaultCase(final boolean defaultCase) {
+        this.defaultCase = defaultCase;
+    }
+
+    public boolean isClassicStyle() {
+        return isClassicStyle;
+    }
+
+    public void setClassicStyle(final boolean classicStyle) {
+        isClassicStyle = classicStyle;
+    }
+
+    @Override
+    public NodeType getNodeType() {
+        return NodeType.UNKNOWN;
+    }
+
+    @Override
+    public <T, R> R acceptVisitor(final IAstVisitor<? super T, ? extends R> visitor, final T data) {
+        return visitor.visitSwitchExpressionArm(this, data);
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public Role<? extends SwitchExpressionArm> getRole() {
+        return (Role<? extends SwitchExpressionArm>) super.getRole();
+    }
+
+    @Override
+    public boolean matches(final INode other, final Match match) {
+        if (other instanceof SwitchExpressionArm) {
+            final SwitchExpressionArm otherSection = (SwitchExpressionArm) other;
+
+            return !otherSection.isNull() &&
+                   getValues().matches(otherSection.getValues(), match) &&
+                   getStatements().matches(otherSection.getStatements(), match);
+        }
+
+        return false;
+    }
+
+    private static final class NullSwitchExpressionArm extends SwitchExpressionArm {
+        public NullSwitchExpressionArm() {
+        }
+
+        @Override
+        public final boolean isNull() {
+            return true;
+        }
+
+        @Override
+        public <T, R> R acceptVisitor(final IAstVisitor<? super T, ? extends R> visitor, final T data) {
+            return null;
+        }
+
+        @Override
+        public boolean matches(final INode other, final Match match) {
+            return other == null || other.isNull();
+        }
+    }
+}
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchSection.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchSection.java
index 7895b016..289ab77e 100644
--- a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchSection.java
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/SwitchSection.java
@@ -21,14 +21,14 @@
 import com.strobel.decompiler.patterns.Role;
 
 public class SwitchSection extends AstNode {
-    public final static Role<CaseLabel> CaseLabelRole = new Role<>("CaseLabel", CaseLabel.class);
+    public final static Role<CaseLabel> CASE_LABEL_ROLE = new Role<>("CaseLabel", CaseLabel.class);
 
     public final AstNodeCollection<Statement> getStatements() {
         return getChildrenByRole(Roles.EMBEDDED_STATEMENT);
     }
 
     public final AstNodeCollection<CaseLabel> getCaseLabels() {
-        return getChildrenByRole(CaseLabelRole);
+        return getChildrenByRole(CASE_LABEL_ROLE);
     }
 
     @Override
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/transforms/RewriteSwitchExpressionsTransform.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/transforms/RewriteSwitchExpressionsTransform.java
new file mode 100644
index 00000000..faf0b743
--- /dev/null
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/transforms/RewriteSwitchExpressionsTransform.java
@@ -0,0 +1,374 @@
+package com.strobel.decompiler.languages.java.ast.transforms;
+
+import com.strobel.core.Predicate;
+import com.strobel.core.StringUtilities;
+import com.strobel.decompiler.DecompilerContext;
+import com.strobel.decompiler.ast.DefaultMap;
+import com.strobel.decompiler.ast.Variable;
+import com.strobel.decompiler.languages.java.analysis.ControlFlowEdge;
+import com.strobel.decompiler.languages.java.analysis.ControlFlowGraphBuilder;
+import com.strobel.decompiler.languages.java.analysis.ControlFlowNode;
+import com.strobel.decompiler.languages.java.analysis.ControlFlowNodeType;
+import com.strobel.decompiler.languages.java.analysis.Correlator;
+import com.strobel.decompiler.languages.java.ast.*;
+import com.strobel.decompiler.patterns.*;
+import com.strobel.functions.Supplier;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static com.strobel.core.CollectionUtilities.*;
+
+public class RewriteSwitchExpressionsTransform extends ContextTrackingVisitor<Void> {
+    private final ExpressionStatement resultAssignment;
+    private final ExpressionStatement firstResultAssignment;
+    private final VariableDeclarationStatement resultDeclaration;
+
+    public RewriteSwitchExpressionsTransform(final DecompilerContext context) {
+        super(context);
+
+        final VariableDeclarationStatement vd = new VariableDeclarationStatement();
+
+        vd.setType(new NamedNode("resultType", new AnyNode()).toType());
+        vd.setAnyModifiers(true);
+
+        vd.getVariables().add(
+            new NamedNode("resultInitializer",
+                          new VariableInitializer(Pattern.ANY_STRING,
+                                                  new NamedNode("initValue",
+                                                                new Choice(new TypedNode(PrimitiveExpression.class),
+                                                                           new NullReferenceExpression())).toExpression())).toVariableInitializer());
+
+        resultDeclaration = vd;
+
+        firstResultAssignment = new ExpressionStatement(
+            new AssignmentExpression(
+                new NamedNode("resultInitializer", new IdentifierExpression(Pattern.ANY_STRING)).toExpression(),
+                new NamedNode("yieldedValue", new TypedNode(Expression.class)).toExpression()
+            )
+        );
+
+        resultAssignment = new ExpressionStatement(
+            new AssignmentExpression(
+                new IdentifierBackReference("resultInitializer").toExpression(),
+                new NamedNode("yieldedValue", new TypedNode(Expression.class)).toExpression()
+            )
+        );
+    }
+
+    @Override
+    public void run(final AstNode compilationUnit) {
+        super.run(compilationUnit);
+    }
+
+    @Override
+    public Void visitSwitchStatement(final SwitchStatement switchNode, final Void data) {
+        super.visitSwitchStatement(switchNode, data);
+
+        final ControlFlowGraphBuilder graphBuilder = new ControlFlowGraphBuilder();
+        final List<ControlFlowNode> nodes = graphBuilder.buildControlFlowGraph(switchNode, new JavaResolver(context));
+        final Set<String> nestedLabels = new HashSet<>();
+
+        switchNode.acceptVisitor(
+            new DepthFirstAstVisitor<Void, Void>() {
+                @Override
+                public Void visitLabelStatement(final LabelStatement node, final Void data) {
+                    nestedLabels.add(node.getLabel());
+                    return super.visitLabelStatement(node, data);
+                }
+
+                @Override
+                public Void visitLabeledStatement(final LabeledStatement node, final Void data) {
+                    nestedLabels.add(node.getLabel());
+                    return super.visitLabeledStatement(node, data);
+                }
+            },
+            null
+        );
+
+        final List<ControlFlowNode> endNodes = toList(
+            where(nodes,
+                  new Predicate<ControlFlowNode>() {
+                      @Override
+                      public boolean test(final ControlFlowNode node) {
+                          final Statement ps = node.getPreviousStatement();
+                          return node.getType() == ControlFlowNodeType.EndNode &&
+                                 ps != switchNode &&
+                                 ps != null &&
+                                 !(ps instanceof BlockStatement) &&
+                                 ps.getParent(SwitchStatement.class) == switchNode;
+                      }
+                  })
+        );
+
+        final SwitchInfo info = new SwitchInfo();
+        final List<BreakStatement> breaks = new ArrayList<>();
+        final List<ThrowStatement> throwStatements = new ArrayList<>();
+        final List<ControlFlowNode> possibleFallThroughNodes = new ArrayList<>();
+
+    outer:
+        for (final ControlFlowNode node : endNodes) {
+            final Statement ps = node.getPreviousStatement();
+
+            if (ps instanceof BreakStatement) {
+                final BreakStatement br = (BreakStatement) ps;
+
+                if (!StringUtilities.isNullOrEmpty(br.getLabel()) && !nestedLabels.contains(br.getLabel())) {
+                    // FAIL: Can't break out of the switch except via `yield` or `throw`.
+                    return null;
+                }
+
+                breaks.add(br);
+            }
+            else if (ps instanceof ThrowStatement) {
+                throwStatements.add((ThrowStatement) ps);
+                continue;
+            }
+            else if (AstNode.isUnconditionalBranch(ps)) {
+                // FAIL: Can't break out of the switch except via `yield` or `throw`.
+                return null;
+            }
+            else {
+                for (final ControlFlowEdge edge : node.getOutgoing()) {
+                    if (edge.getTo() != null && edge.getTo().getOutgoing().size() != 0) {
+                        continue outer;
+                    }
+                }
+                possibleFallThroughNodes.add(node);
+                info.needClassicStyle = true;
+            }
+        }
+
+        IdentifierExpression firstAssignment = null;
+
+        for (final BreakStatement br : breaks) {
+            final Match m = Match.createNew();
+            final Statement bp = br.getPreviousStatement();
+
+            if (firstAssignment != null) {
+                m.add("resultInitializer", firstAssignment);
+                if (!resultAssignment.matches(bp, m)) {
+                    // FAIL: No assignment prior to the break.
+                    return null;
+                }
+            }
+            else if (firstResultAssignment.matches(bp, m)) {
+                firstAssignment = first(m.<IdentifierExpression>get("resultInitializer"));
+                info.resultIdentifier = firstAssignment;
+                info.resultVariable = firstAssignment.getUserData(Keys.VARIABLE);
+
+                if (info.resultVariable == null) {
+                    // FAIL: Need to be able to correlate reads on the variable.
+                    return null;
+                }
+            }
+            else {
+                // FAIL: No assignment prior to the break.
+                return null;
+            }
+
+            final SwitchSection ss = br.getParent(SwitchSection.class);
+
+            if (ss == null || ss.isNull()) {
+                // FAIL: This isn't supposed to be possible.
+                return null;
+            }
+
+            final CaseInfo caseInfo = info.cases.get(ss);
+
+            caseInfo.breaks.add(br);
+            caseInfo.yieldedValues.put(br, first(m.<Expression>get("yieldedValue")));
+
+            caseInfo.isDefault = ss.getCaseLabels().isEmpty() || ss.getCaseLabels().size() == 1 &&
+                                                                 ss.getCaseLabels().firstOrNullObject().getExpression().isNull();
+        }
+
+        for (final ThrowStatement ts : throwStatements) {
+            final SwitchSection ss = ts.getParent(SwitchSection.class);
+
+            if (ss == null || ss.isNull()) {
+                // FAIL: This isn't supposed to be possible.
+                return null;
+            }
+
+            info.cases.get(ss).hasThrows = true;
+        }
+
+        for (final ControlFlowNode fallThrough : possibleFallThroughNodes) {
+            final Statement ps = fallThrough.getPreviousStatement();
+            final SwitchSection ss = ps != null ? ps.getParent(SwitchSection.class) : null;
+            final CaseInfo caseInfo;
+
+            if (ss == null || !(caseInfo = info.cases.get(ss)).breaks.isEmpty()) {
+                // FAIL: Somehow we're breaking out of the switch
+                return null;
+            }
+
+            caseInfo.hasFallThrough = true;
+        }
+
+        boolean fallThroughNeeded = false;
+
+        for (final SwitchSection section : switchNode.getSwitchSections()) {
+            final boolean isEmpty = !info.cases.containsKey(section);
+            final CaseInfo caseInfo = info.cases.get(section);
+
+            if (isEmpty) {
+                caseInfo.isEmpty = true;
+            }
+
+            if (isEmpty || caseInfo.hasFallThrough) {
+                fallThroughNeeded = true;
+            }
+            else if (caseInfo.hasThrows || !caseInfo.yieldedValues.isEmpty()) {
+                fallThroughNeeded = false;
+            }
+        }
+
+        if (fallThroughNeeded) {
+            // FAIL: Some section didn't produce a value, and none of its successors
+            //       produced a value either.
+            return null;
+        }
+
+        final boolean areCorrelated = Correlator.areCorrelated(info.resultIdentifier.clone(), switchNode);
+
+        if (areCorrelated) {
+            // FAIL: Our target variable gets read somewhere inside the switch.
+            return null;
+        }
+
+        rewrite(switchNode, info);
+        return null;
+    }
+
+    private void rewrite(final SwitchStatement node, final SwitchInfo info) {
+        final AstNodeCollection<SwitchSection> sections = node.getSwitchSections();
+        final SwitchExpression se = new SwitchExpression();
+
+        for (final SwitchSection ss : sections) {
+            final CaseInfo caseInfo = info.cases.get(ss);
+            final AstNodeCollection<Statement> statements = ss.getStatements();
+            final Statement first = statements.firstOrNullObject();
+            final SwitchExpressionArm arm = new SwitchExpressionArm();
+            final BlockStatement b;
+
+            if (statements.size() == 1 && first instanceof BlockStatement) {
+                b = (BlockStatement) first;
+                b.remove();
+            }
+            else {
+                b = new BlockStatement();
+
+                for (final Statement statement : statements) {
+                    statement.remove();
+                    b.add(statement);
+                }
+            }
+
+            for (final CaseLabel label : ss.getCaseLabels()) {
+                final Expression value = label.getExpression();
+
+                if (value.isNull()) {
+                    arm.setDefaultCase(true);
+                }
+                else {
+                    value.remove();
+                    arm.getValues().add(value);
+                }
+            }
+
+            for (final BreakStatement br : caseInfo.breaks) {
+                final Expression value = caseInfo.yieldedValues.get(br);
+                final Statement parentStatement = value.getParent(Statement.class);
+
+                assert parentStatement != null;
+
+                parentStatement.remove();
+                value.remove();
+
+                br.setYield(true);
+                br.setValue(value);
+            }
+
+            final Statement newFirst;
+
+            if (!info.needClassicStyle &&
+                b.getStatements().size() == 1 &&
+                !(newFirst = b.getStatements().firstOrNullObject()).isNull()) {
+
+                if (newFirst instanceof BreakStatement) {
+                    final Expression value = ((BreakStatement) newFirst).getValue();
+                    value.remove();
+                    arm.getStatements().add(new ExpressionStatement(value));
+                }
+                else if (newFirst instanceof ThrowStatement) {
+                    newFirst.remove();
+                    arm.getStatements().add(newFirst);
+                }
+                else {
+                    arm.getStatements().add(b);
+                }
+            }
+            else {
+                arm.getStatements().add(b);
+            }
+
+            arm.setClassicStyle(info.needClassicStyle);
+            se.getArms().add(arm);
+        }
+
+        final Expression governingExpression = node.getExpression();
+
+        governingExpression.remove();
+        se.setGoverningExpression(governingExpression);
+
+        final Statement ps = node.getPreviousStatement();
+        final Match m = ps != null ? resultDeclaration.match(ps) : Match.failure();
+
+        final VariableInitializer vi;
+
+        if (m.success() &&
+            info.resultVariable == (vi = first(m.<VariableInitializer>get("resultInitializer"))).getUserData(Keys.VARIABLE)) {
+
+            vi.getInitializer().replaceWith(se);
+            node.remove();
+        }
+        else {
+            node.replaceWith(new ExpressionStatement(new AssignmentExpression(info.resultIdentifier.clone(), se)));
+        }
+    }
+
+    final static class CaseInfo {
+        final List<BreakStatement> breaks = new ArrayList<>();
+        final Map<BreakStatement, Expression> yieldedValues = new HashMap<>();
+
+        boolean isEmpty;
+        boolean isDefault;
+        boolean hasThrows;
+        boolean hasFallThrough;
+
+        CaseInfo() {
+        }
+    }
+
+    final static class SwitchInfo {
+        IdentifierExpression resultIdentifier;
+        Variable resultVariable;
+
+        final Map<SwitchSection, CaseInfo> cases = new DefaultMap<>(
+            new Supplier<CaseInfo>() {
+                @Override
+                public CaseInfo get() {
+                    return new CaseInfo();
+                }
+            }
+        );
+
+        boolean needClassicStyle;
+    }
+}
diff --git a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/transforms/TransformationPipeline.java b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/transforms/TransformationPipeline.java
index 314208dc..b994cf64 100644
--- a/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/transforms/TransformationPipeline.java
+++ b/Procyon.CompilerTools/src/main/java/com/strobel/decompiler/languages/java/ast/transforms/TransformationPipeline.java
@@ -55,6 +55,7 @@ public static IAstTransform[] createPipeline(final DecompilerContext context) {
             new RewriteInnerClassConstructorCalls(context),
             new RemoveRedundantInitializersTransform(context),
             new FlattenElseIfStatementsTransform(context),
+            new RewriteSwitchExpressionsTransform(context),
             new FlattenSwitchBlocksTransform(context),
             new IntroduceInitializersTransform(context),
             new MarkReferencedSyntheticsTransform(context),